A General Framework for Inference-time Scaling and Steering of Diffusion Models

Diffusion models produce impressive results in modalities ranging from images and video to protein design and text. However, generating samples with user-specified properties remains a challenge. Recent research proposes fine-tuning models to maximize rewards that capture desired properties, but these methods require expensive training and are prone to mode collapse. In this work, we present Feynman-Kac (FK) steering, an inference-time framework for steering diffusion models with reward functions. FK steering works by sampling a system of multiple interacting diffusion processes, called particles, and resampling particles at intermediate steps based on scores computed using functions called potentials. Potentials are defined using rewards for intermediate states and are selected such that a high value indicates that the particle will yield a high-reward sample. We explore various choices of potentials, intermediate rewards, and samplers. We evaluate FK steering on text-to-image and text diffusion models. For steering text-to-image models with a human preference reward, we find that FK steering a 0.8B parameter model outperforms a 2.6B parameter fine-tuned model on prompt fidelity, with faster sampling and no training. For steering text diffusion models with rewards for text quality and specific text attributes, we find that FK steering generates lower perplexity, more linguistically acceptable outputs and enables gradient-free control of attributes like toxicity. Our results demonstrate that inference-time scaling and steering of diffusion models - even with off-the-shelf rewards - can provide significant sample quality gains and controllability benefits. Code is available at https://github.com/zacharyhorvitz/Fk-Diffusion-Steering .

한 줄 요약

Feynman-Kac (FK) steering을 통해 추론 시 확산 모델을 보상 함수로 제어하고, 훈련 없이도 샘플 품질과 제어력을 향상시킨다.

핵심 기여도

FK steering은 훈련 없이 추론 시 확산 모델을 제어하는 새로운 프레임워크로, 0.8B 파라미터 모델이 2.6B 파라미터 미세조정 모델보다 샘플 품질에서 우수함.
텍스트 확산 모델에서 토스티시티와 같은 속성을 기반으로 샘플을 제어할 수 있으며, 퍼플렉시티가 낮아 더 자연스러운 결과 생성.
다양한 보상 함수와 중간 보상, 포텐셜 함수를 실험적으로 탐색하여 사용자 정의 제어가 가능.
4개의 파티클만 사용해 기존 best-of-n 샘플링보다 높은 성능을 보임.

핵심 아이디어

FK steering은 확산 모델의 추론 단계에서 **Feynman-Kac 상호작용 파티클 시스템**을 활용하여 샘플 품질과 제어력을 동시에 향상시키는 새로운 접근법이다. 이는 **보상 함수**를 기반으로 샘플을 가중치를 주어 선택적으로 생성하는 방식이다.

기존 방법은 미세조정(fine-tuning)이나 기울기 기반 가이드라인을 사용했지만, 이들은 훈련 비용이 높거나 보상 함수의 미분 가능성에 의존한다. FK steering은 **보상 함수의 미분 가능성 여부와 관계없이 작동**하며, **중간 단계에서 파티클을 재샘플링**함으로써 높은 보상 샘플을 생성한다.

핵심 아이디어는 **포텐셜 함수**(potential)를 사용하여 각 파티클의 중간 상태를 평가하고, 높은 포텐셜 값을 가진 파티클을 확장하는 것이다. 이는 **중간 보상**(intermediate reward)을 기반으로 정의되며, 샘플의 최종 품질을 예측하는 역할을 한다.

기술적 접근법

**FK steering 프레임워크**는 여러 확산 프로세스(파티클)를 생성하고, 중간 단계에서 **포텐셜 함수**를 사용해 파티클을 재샘플링한다.
**포텐셜 함수**는 중간 보상(예: 텍스트 품질, 이미지 품질)을 기반으로 정의되며, 높은 포텐셜은 높은 보상 샘플을 예측한다.
실험에서는 **ImageReward**, **GPT2**, **trigram 언어 모델**, **언어 수용성 분류기** 등 다양한 보상 함수를 사용.
**파티클 수**는 4개로 설정해도 기존 best-of-n 샘플링보다 우수한 성능을 보임.
**제안 생성기**(proposal generator)로는 기본 확산 모델을 사용하며, 이는 훈련 없이도 작동 가능.

주요 결과

**GenEval 벤치마크**에서 FK steering을 적용한 0.8B 파라미터 모델은 2.6B 파라미터 미세조정 모델보다 **프롬프트 일치도(prompt fidelity)**에서 우수함.
텍스트 확산 모델에서 FK steering은 **퍼플렉시티(perplexity)**가 낮고, **언어 수용성**(linguistic acceptability)이 높은 결과를 생성.
**1%**의 샘플만 발생하는 토스티시티(toxicity) 속성도 **기울기 없이 제어** 가능.
**4개의 파티클**만 사용해도 기존 best-of-n 샘플링을 초과하는 성능을 보임.

의의 및 한계

FK steering은 추론 시 확산 모델을 제어하는 새로운 패러다임으로, **훈련 없이도 샘플 품질과 제어력을 동시에 향상**시킨다. 특히, **작은 모델**(0.8B)이 **큰 모델**(2.6B)보다 성능이 우수한 점은 계산 효율성 측면에서 중요한 의의를 가진다. 또한, **보상 함수의 유연한 선택**이 가능해 다양한 사용 사례에 적용 가능하다.

하지만, FK steering의 성능은 **강력한 보상 함수**에 의존하며, 보상 함수가 부족하거나 불완전하면 제어력이 떨어질 수 있다. 또한, **파티클 수를 늘릴수록 계산 비용이 증가**하므로, 성능-비용 트레이드오프를 고려해야 한다.

실용적 활용

FK steering은 **텍스트-이미지 생성**, **텍스트 생성**, **단백질 설계** 등 다양한 생성 모델에 적용 가능하다. 특히, **사용자 선호도**(human preference)나 **속성**(toxicity, 클래스) 기반 샘플 생성이 필요한 산업 및 연구 분야에서 유용하다. 예를 들어, **광고 이미지 생성**, **콘텐츠 필터링**, **생물학적 설계** 등에서 실용적 활용이 기대된다.