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

• URL: http://arxiv.org/abs/2501.06848v3
• Date: Thu, 16 Jan 2025 03:18:14 GMT
• Title: A General Framework for Inference-time Scaling and Steering of Diffusion Models
• Authors: Raghav Singhal, Zachary Horvitz, Ryan Teehan, Mengye Ren, Zhou Yu, Kathleen McKeown, Rajesh Ranganath,
• Abstract summary: Feynman Kac (FK) steering is 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. We evaluate FK steering on text-to-image and text diffusion models.
• Score: 45.27503904127995
• License:
• Abstract: 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 propose 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 .

Related papers

• Steering Masked Discrete Diffusion Models via Discrete Denoising Posterior Prediction [88.65168366064061]
  We introduce Discrete Denoising Posterior Prediction (DDPP), a novel framework that casts the task of steering pre-trained MDMs as a problem of probabilistic inference. Our framework leads to a family of three novel objectives that are all simulation-free, and thus scalable. We substantiate our designs via wet-lab validation, where we observe transient expression of reward-optimized protein sequences.
  arXiv  Detail & Related papers  (2024-10-10T17:18:30Z)
• Feedback Efficient Online Fine-Tuning of Diffusion Models [52.170384048274364]
  We propose a novel reinforcement learning procedure that efficiently explores on the manifold of feasible samples. We present a theoretical analysis providing a regret guarantee, as well as empirical validation across three domains.
  arXiv  Detail & Related papers  (2024-02-26T07:24:32Z)
• Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control [54.132297393662654]
  Diffusion models excel at capturing complex data distributions, such as those of natural images and proteins. While diffusion models are trained to represent the distribution in the training dataset, we often are more concerned with other properties, such as the aesthetic quality of the generated images. We present theoretical and empirical evidence that demonstrates our framework is capable of efficiently generating diverse samples with high genuine rewards.
  arXiv  Detail & Related papers  (2024-02-23T08:54:42Z)
• Value function estimation using conditional diffusion models for control [62.27184818047923]
  We propose a simple algorithm called Diffused Value Function (DVF) It learns a joint multi-step model of the environment-robot interaction dynamics using a diffusion model. We show how DVF can be used to efficiently capture the state visitation measure for multiple controllers.
  arXiv  Detail & Related papers  (2023-06-09T18:40:55Z)
• Extracting Reward Functions from Diffusion Models [7.834479563217133]
  Decision-making diffusion models can be trained on lower-quality data, and then be steered with a reward function to generate near-optimal trajectories. We consider the problem of extracting a reward function by comparing a decision-making diffusion model that models low-reward behavior and one that models high-reward behavior. We show that our approach generalizes beyond sequential decision-making by learning a reward-like function from two large-scale image generation diffusion models.
  arXiv  Detail & Related papers  (2023-06-01T17:59:12Z)
• Diff-Instruct: A Universal Approach for Transferring Knowledge From Pre-trained Diffusion Models [77.83923746319498]
  We propose a framework called Diff-Instruct to instruct the training of arbitrary generative models. We show that Diff-Instruct results in state-of-the-art single-step diffusion-based models. Experiments on refining GAN models show that the Diff-Instruct can consistently improve the pre-trained generators of GAN models.
  arXiv  Detail & Related papers  (2023-05-29T04:22:57Z)
• Towards Controllable Diffusion Models via Reward-Guided Exploration [15.857464051475294]
  We propose a novel framework that guides the training-phase of diffusion models via reinforcement learning (RL) RL enables calculating policy gradients via samples from a pay-off distribution proportional to exponential scaled rewards, rather than from policies themselves. Experiments on 3D shape and molecule generation tasks show significant improvements over existing conditional diffusion models.
  arXiv  Detail & Related papers  (2023-04-14T13:51:26Z)
• Understanding Diffusion Models: A Unified Perspective [0.0]
  Diffusion models have shown incredible capabilities as generative models. We review, demystify, and unify the understanding of diffusion models across both variational and score-based perspectives.
  arXiv  Detail & Related papers  (2022-08-25T09:55:25Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.