PaperlayerView exact revision on arXiv

Source-linked AI summary

Inference-Time Alignment in Diffusion Models with Reward-Guided Generation: Tutorial and Review

Masatoshi Uehara, Yulai Zhao, Chenyu Wang, Xiner Li, Aviv Regev, Sergey Levine, Tommaso Biancalani

arXiv:2501.09685v2cs.AIcs.LGq-bio.QMstat.ML

TL;DR

Diffusion models need inference-time methods to optimize downstream rewards such as protein affinity or stability while retaining naturalistic generation. This tutorial unifies guidance techniques as approximations to soft-optimal denoising policies and explains their formal grounding and practical scope.

  • Problem

    Applications require diffusion models to generate natural designs while maximizing downstream rewards such as protein binding affinity or stability.

  • Method

    The tutorial reviews inference-time guidance methods through a unified soft-optimal-policy framework and surveys their applications, extensions, and related search-based approaches.

  • Results

    The reviewed methods approximate soft-optimal denoising policies, while classifier guidance can formally target distributions proportional to exp(r(x))p_pre(x).

  • Takeaways & Limitations

    Sequentially sampling from these soft-optimal policies can generate natural-like designs with high functionality during inference.

  • Takeaways & Limitations

    Derivative-based guidance lacks formal guarantees in discrete spaces because its Taylor expansion is technically invalid there.

Abstract

from arXiv · show

This tutorial provides an in-depth guide on inference-time guidance and alignment methods for optimizing downstream reward functions in diffusion models. While diffusion models are renowned for their generative modeling capabilities, practical applications in fields such as biology often require sample generation that maximizes specific metrics (e.g., stability, affinity in proteins, closeness to target structures). In these scenarios, diffusion models can be adapted not only to generate realistic samples but also to explicitly maximize desired measures at inference time without fine-tuning. This tutorial explores the foundational aspects of such inference-time algorithms. We review these methods from a unified perspective, demonstrating that current techniques -- such as Sequential Monte Carlo (SMC)-based guidance, value-based sampling, and classifier guidance -- aim to approximate soft optimal denoising processes (a.k.a. policies in RL) that combine pre-trained denoising processes with value functions serving as look-ahead functions that predict from intermediate states to terminal rewards. Within this framework, we present several novel algorithms not yet covered in the literature. Furthermore, we discuss (1) fine-tuning methods combined with inference-time techniques, (2) inference-time algorithms based on search algorithms such as Monte Carlo tree search, which have received limited attention in current research, and (3) connections between inference-time algorithms in language models and diffusion models. The code of this tutorial on protein design is available at https://github.com/masa-ue/AlignInversePro

Introduction

The tutorial frames inference-time alignment as reward-guided controlled generation that optimizes downstream functionality without diffusion-model fine-tuning. It unifies existing methods as approximations to soft optimal denoising policies combining pretrained processes with value-based look-ahead, while covering extensions to search, editing, language models, and fine-tuning.

  • Introduction: Inference-time guidance optimizes downstream rewards such as protein binding affinity or stability after conditional pretraining, without directly fine-tuning the diffusion model.The same framework treats inference-time conditioning on target properties as reward maximization, often using classifiers.
  • Inference-Time Techniques vs. Post-Training: Compared with post-training, inference-time methods are often fine-tuning-free or training-free, can improve rewards by increasing inference compute, and remain useful alongside or after fine-tuning.They can augment data, serve as teacher policies for distillation, and directly exploit accurate reward feedback that may be degraded when converted into training data.
  • Critical Considerations for Choosing Inference-Time Techniques: Technique selection depends on compute and memory, reward type and efficiency, and whether reward feedback is differentiable, especially for simulation- or descriptor-based molecular objectives.The tutorial distinguishes classifier-based conditioning from regressor-based alignment and notes that nondifferentiable feedback motivates gradient-free methods.
  • Summary: Inference-time methods approximate soft optimal denoising policies that combine pretrained generation with value functions predicting terminal rewards, producing natural and functional outputs.The target distribution balances the pretrained distribution with a reward term, while methods differ in how they approximate the corresponding denoising process.
  • Summary: The tutorial reviews gradient-free SMC and value-based importance sampling for nondifferentiable rewards, alongside differentiable classifier guidance for continuous and discrete diffusion.It also explains their integration and formalizes continuous classifier guidance through Doob’s transform.
  • Summary: Beyond core guidance, the tutorial covers search algorithms such as MCTS, sequential refinement for editing endogenous designs, and inference-time methods for autoregressive and masked language models.These extensions address advanced search, constrained protein editing, and similarities or differences across language-model and diffusion frameworks.
  • Summary: The tutorial describes distillation-based fine-tuning of diffusion models, noting higher inference costs for pure inference-time methods and connections between these methods and RL-based fine-tuning.It also reviews protein-design approaches, including walk-jump sampling and hallucination methods such as genetic algorithms and MCMC-based approaches.

Application in Computational Protein Design

The tutorial illustrates inference-time alignment for computational protein design by pairing foundational diffusion models with reward functions that encode desired protein properties. It surveys sequence, structure-based, and hybrid generators, alongside reward construction from experimental data, learned features, or physics-based models.

  • Overview: Inference-time protein design requires selecting a pre-trained diffusion model and defining a reward function, as illustrated in Figure 1.The tutorial’s implementation code is available for this protein-design example.
  • Pre-Trained Foundational Diffusion Models for Protein Sequences: Purely sequence-based models train discrete diffusion processes on protein sequence datasets without structural data, exemplified by EvoDiff.
  • Pre-Trained Foundational Diffusion Models for Protein Sequences: Structure-based pipelines first generate backbones with RFdiffusion or Chroma, then perform inverse folding with ProteinMPNN.
  • Pre-Trained Foundational Diffusion Models for Protein Sequences: Hybrid models jointly generate sequences and structures at each step, including MultifLow (Campbell et al., 2024), DPLM-2 (Wang et al., 2024), Protein Generator, and ESM3.Foundational models may also condition generation on secondary structures, motifs, domains, or symmetry.
  • Reward Models (Mapping Protein Sequences to Functionality): Reward models target binding affinity, stability, solubility, or rigidity using experimental sequence-property data supplemented by learned representations, predicted structures, or biophysical features such as ∆G.When experimental data are unavailable, rewards can use AlphaFold3, Rosetta, or AutoDock.

1 Preliminaries

This section introduces diffusion models as sequential denoising processes and formalizes inference-time conditioning, inverse problems, and reward maximization while preserving sample naturalness. It also distinguishes these objectives by their reward functions and highlights practical challenges in reward computation and multi-objective optimization.

  • 1.1 Diffusion Models: Diffusion models learn reverse denoising transitions that sequentially transform noise into samples approximating the training-data distribution.The forward process defines noising from t = 0 to T, while the learned reverse process is sampled sequentially from t = T to t = 0.
  • 1.2 Objectives: Conditioning, Inverse Problems, Reward Maximization: Inference-time objectives seek high-reward designs while preserving naturalness, represented by reward optimization combined with a term favoring the pretrained sample distribution.Rewards may target objectives such as protein binding affinity, while the naturalness term keeps generated samples aligned with the learned data distribution.
  • 1.2.1 Conditioning: Conditioning uses classifier-based log-likelihood rewards log p(y | x) to generate samples from a desired conditional distribution p(x | y).The classifier is typically learned from paired (x, y) data; protein-function classifiers provide an example.
  • 1.2.2 Inverse Problems: Inverse problems use known observation models and likelihood-based rewards, such as Gaussian-noise reconstruction error, for tasks including image inpainting and protein motif scaffolding.The measurement operator A maps designs to observations y, and Gaussian noise yields an explicit reward based on ∥y − A(x)∥2.
  • 1.2.3 Reward Maximization (i.e., Alignment): Alignment problems optimize downstream regressor-based rewards, including protein binding affinity, immunogenicity, designability, self-consistency, and stability.Rewards may be predefined from simulations or cheminformatics tools, or learned from experimental data.
  • 1.2.3 Reward Maximization (i.e., Alignment): Alignment must balance reward accuracy against computational cost and reconcile conflicting objectives such as affinity, specificity, lipophilicity, toxicity, and designability.High-resolution simulations can be prohibitive, while multi-objective molecular design requires trade-offs among competing rewards.

2 Foundations of Inference-Time Controlled Generation

This section formulates inference-time controlled generation as approximation of a soft-optimal denoising policy that combines pretrained diffusion processes with reward-based look-ahead values. It then presents value-function approximation strategies used by later inference-time methods.

  • 2.1 Soft Optimal Policies (Denoising Processes) in Diffusion Models: Soft-optimal denoising policies retain pretrained-generation characteristics while steering samples toward high-reward, natural-like designs through entropy-regularized value functions.The soft value function predicts future terminal rewards from intermediate states and guides sequential inference.
  • 2.1 Soft Optimal Policies (Denoising Processes) in Diffusion Models: All subsequent guidance methods approximate this policy, using derivative-free approaches in Section 3 and value-function gradients in Sections 4 and 5.Exact computation is difficult because value functions are unavailable and the action space makes normalization and value evaluation computationally expensive.
  • 2.2 Approximating Soft Value Functions: Inference-time techniques require approximating soft value functions because the target policy is expressed through pretrained policies multiplied by value functions.These approximations can be integrated into the inference-time methods discussed later.
  • 2.2.1 Posterior Mean Approximation: Posterior mean approximation estimates intermediate value by decoding x_t to x̂_0 and evaluating r(x̂_0), as used in DPS (Chung et al., 2022), universal guidance, and reconstruction guidance (Ho et al., 2022).The method uses the pretrained model’s decoder from x_t to x_0.
  • 2.2.2 Monte Carlo Regression: Monte Carlo regression trains a value-function regressor by replacing the soft-value expectation with an empirical objective under a roll-in distribution.When pretrained-model-induced distributions provide the roll-in distribution, this yields the procedure summarized in Algorithm 2.
  • 2.2.3 Soft Q-Learning: Soft Q-learning estimates value functions by recursively applying regression through fitted Q-iteration, using roll-in distributions to approximate the soft-Bellman recursion.The resulting procedure is summarized in Algorithm 3.

3 Derivative-Free Guidance

Derivative-free guidance addresses non-differentiable reward feedback common in molecular design, using SMC-based guidance and value-based importance sampling without requiring differentiable reward models. The section further develops nested-SMC, beam search, and proposal-distribution choices for reward-guided diffusion.

  • 3 Derivative-Free Guidance: Derivative-free methods are useful when rewards come from black-box simulations, non-differentiable molecular descriptors, or non-differentiable models such as graph neural networks and XGBoost.These settings motivate approaches that do not rely on differentiable reward models.
  • 3.1 SMC-Based Guidance: SMC-based guidance combines diffusion models with particle filtering, using importance weighting and effective-sample-size-triggered resampling to approximate optimal policies, but may collapse samples and fail to ensure diversity.Its global batch interactions eliminate inferior particles and concentrate computation on promising samples.
  • 3.2 Value-Based Importance Sampling: For reward maximization with small α, SVDD is preferable because it independently resamples within each sample and avoids the mode-collapse risk associated with global SMC interactions.For conditioning with moderate α, SMC can instead be advantageous because global interaction removes inferior samples.
  • 3.3 Nested-SMC-Based Guidance: Nested-SMC guidance combines SVDD-like local sampling with SMC-like global resampling, enabling suboptimal samples to be eliminated while retaining both methods’ complementary strengths.The method is presented as a hybrid nested-importance-sampling strategy.
  • 3.4 Beam Search for Reward Maximization: When α = 0, SVDD becomes beam search that expands multiple proposal nodes and selects the highest-value node using soft value functions as look-ahead mechanisms.This interpretation assumes perfect soft-value access; approximation errors or deeper-search benefits motivate considering MCTS.
  • 3.5 Selecting Proposal Distributions: Proposal distributions can use pre-trained diffusion policies, derivative-based guidance with differentiable value models, or fine-tuned policies, with differentiable surrogates sometimes preserving useful reward signals.Fine-tuned policies can serve as enhanced proposals when combined with distillation and inference-time techniques.

4 Derivative-Based Guidance in Continuous Diffusion Models · 4.1 Intuitive Derivation of Classifier Guidance · 4.2 Derivative-Free Guidance Versus Classifier Guidance

This section presents classifier guidance as a continuous-time, derivative-based approximation to the soft-optimal policy and extends it to Riemannian diffusion models. It contrasts this approach with derivative-free guidance, emphasizing their assumptions and possible combination.

  • 4 Derivative-Based Guidance in Continuous Diffusion Models: The method extends to Riemannian diffusion models by replacing Euclidean gradients with Riemannian gradients.This extension is intended for settings such as protein structure generation.
  • 4 Derivative-Based Guidance in Continuous Diffusion Models: Classifier guidance adds value-function gradients at inference time to approximate the soft-optimal policy when discretization errors are negligible (Dhariwal and Nichol, 2021; Song et al., 2021).The continuous-time formulation uses Doob’s transform.
  • 4.1 Intuitive Derivation of Classifier Guidance: Classifier guidance derives a Gaussian policy that approximates the soft-optimal policy by applying a Taylor expansion to an otherwise difficult-to-sample target policy.The derivation assumes a small diffusion step and approximates the policy up to a normalizing constant.
  • 4.1 Intuitive Derivation of Classifier Guidance: The resulting continuous-diffusion algorithm incorporates the gradient of the value function during inference using pretrained diffusion models and differentiable soft value functions (Dhariwal and Nichol, 2021).The complete procedure is summarized as Algorithm 8.
  • 4.2 Derivative-Free Guidance Versus Classifier Guidance: Classifier guidance critically depends on accurate differentiable value functions, making it effective for inpainting but potentially unsuitable for molecular design tasks with nondifferentiable rewards.Derivative-free guidance avoids requiring differentiable rewards or value-function models.
  • 4.2 Derivative-Free Guidance Versus Classifier Guidance: Classifier and derivative-free guidance can be combined by using classifier guidance as a proposal distribution while SMC-based guidance or value-based sampling uses more precise nondifferentiable value functions.This remains possible even when the differentiable models used for classifier guidance are imperfect.

4.3 Continuous-Time Formalization via Doob transform

This section formalizes classifier guidance in continuous time by combining diffusion-model time reversal with a Doob transform. The resulting controlled SDE samples from a reward-tilted version of the pre-trained distribution, subject to an initial-distribution adjustment when the model’s initial state is stochastic.

  • 4.3.1 Preparation: The continuous-time diffusion framework defines a forward SDE and learns its time-reversal SDE by estimating score functions.The forward process approaches N(0, I) as T approaches ∞, while the time-reversal SDE preserves marginal distributions.
  • 4.3.1 Preparation: Starting from N(0, I) and applying the learned time-reversal SDE recovers the data distribution when T is sufficiently large.The analysis assumes a fixed pre-trained score model s(z_t, T −t; θ_pre) that predicts ∇log q_T−t(z_t).
  • 4.3.2 Doob Transform: The Doob transform augments the pre-trained drift with ∇log v_t(z_t), where v_t is the exponentiated-reward look-ahead value under the pre-trained process.Specifically, v_t(·) = log E_θpre[exp(r(z_T))|z_t = ·] in the theorem’s notation.
  • 4.3.2 Doob Transform: The transformed SDE induces a target distribution proportional to exp(r(x))p_pre(x).Thus, classifier guidance with standard Euler-Maruyama discretization and α = 1 samples from this reward-tilted target distribution.
  • 4.3.2 Doob Transform: The continuous-time value v_t is the analogue of the value function used in the discrete-time guidance formulation.The theorem therefore provides a formal continuous-time derivation of classifier guidance rather than treating the guidance update heuristically.
  • 4.3.2 Doob Transform: The Doob-transform characterization connects classifier guidance to entropy-regularized optimal control in reinforcement-learning-based fine-tuning.These connections are attributed to Zhao et al. (2024), Denker et al. (2024), and Uehara et al. (2024).
  • 4.3.2 Doob Transform: When the pre-trained model has a stochastic initial distribution, the initial law must instead be proportional to exp(v_0(x))p_ini(x).This adjustment is required to preserve the formal target-distribution result.

4.4 Guidance in Riemannian Diffusion Models

This section extends inference-time guidance from Euclidean spaces to Riemannian manifolds, focusing on SO(3) because SE(3) models protein-backbone conformations (Yim et al., 2023; Watson et al., 2023). It introduces manifold geometry and adapts classifier guidance by adding Riemannian gradients before geodesic updates.

  • 4.4.1 Primer: Riemannian Manifolds: Riemannian manifolds locally resemble Euclidean space and have tangent spaces with metrics; SO(3) has intrinsic dimension 3, is embedded in R9, and has skew-symmetric tangent matrices.
  • 4.4.2 Classifier Guidance in Riemannian Diffusion Models: Riemannian diffusion updates first compute a tangent-space velocity and then move along its induced geodesic, generalizing the Euclidean update xt + velt.Noise is sampled from a distribution on the manifold, while the pretrained model is formulated as an SDE with Brownian motion on the manifold.
  • 4.4.2 Classifier Guidance in Riemannian Diffusion Models: Classifier guidance adds the Riemannian gradient of the soft value function to the pretrained velocity at every inference step, then outputs x0 after geodesic transitions.For SO(3), the Riemannian gradient is computed using the manifold-specific gradient expression; the method can be formalized with Doob’s theorem for Riemannian manifolds.

5 Derivative-Based Guidance in Discrete Diffusion Models

This section shows that discrete diffusion models admit polynomial-time exact sampling from reward-optimal policies, motivates derivative-based guidance as a cheaper approximation, and formalizes classifier guidance through continuous-time Doob transforms. The formal result identifies the guided process as sampling from the reward-reweighted pretrained distribution, while derivative approximations lack discrete-space guarantees.

  • 5.1 Exact Sampling in Discrete Diffusion Models: Discrete diffusion models have effective action spaces scaling as LK, enabling exact optimal-policy sampling with polynomial-time value-function computation rather than exponential computation.The same LK cost can remain prohibitive in practice, motivating SMC-based, value-based, and derivative-based approximations.
  • 5.2 Derivative-Based Guidance in Discrete Diffusion Models: Derivative-based guidance approximates classifier guidance by Taylor-expanding a soft value function in an L×K one-hot representation, as proposed by Nisonoff et al. (2024).The approximation ignores higher-order terms and uses value-function derivatives to reduce the LK computation overhead.
  • 5.2 Derivative-Based Guidance in Discrete Diffusion Models: Derivative-based guidance has no formal guarantees in discrete spaces because formal derivatives do not exist there, making the Taylor expansion technically invalid.This limitation distinguishes the practical approximation from the continuous-domain setting.
  • 5.3.1 Preparation: In continuous-time discrete diffusion, CTMC ratio matching replaces score matching, and the pretrained process is represented through forward and time-reversed generators.The time-reversal process can sample the data distribution when the marginal density ratio is learned, assuming pretraining is complete.
  • 5.3.2 Doob Transform in CTMC: The Doob-transform theorem defines vt(·) as the log pretrained expectation of exp(r(zT)) and constructs a CTMC whose distribution is proportional to exp(r(x))ppre(x).This provides the continuous-time formalization of classifier guidance in discrete diffusion models.
  • 5.3.2 Doob Transform in CTMC: Using Euler–Maruyama discretization, the continuous-time value function yields Algorithm 10, which samples from the target reward-reweighted distribution.Wang et al. (2024) additionally proves that this CTMC maximizes the entropy-regularized reward objective in RL-based fine-tuning.

6 Tree Search Algorithms for Alignment

Tree search can align diffusion-model generation by exploiting the tree structure induced by pretrained denoising processes and using value functions to maximize rewards. The section focuses on controlling the resulting tree’s width and depth and improving leaf evaluation through pretrained-model rollouts, while noting computational trade-offs and limited prior diffusion-model research.

  • 6 Tree Search Algorithms for Alignment: Pre-trained diffusion models induce a tree structure that MCTS can search to maximize rewards using value functions.This direction has been less explored for diffusion models than language models, despite MCTS’s success in molecular generation.
  • 6.1 Defining the Search Tree: Diffusion search trees have substantial depth and width, making exhaustive exploration computationally prohibitive in continuous and masked discrete settings.Depth typically follows 50–1000 discretization levels, while masked discrete diffusion has width LK for token length L and vocabulary size K.
  • 6.1 Defining the Search Tree: Search algorithms constrain tree growth by sampling a limited number of nodes from pretrained models during expansion.This strategy limits both width and depth from the current inference state and parallels subsampling methods used in AlphaZero variants and language models.
  • 6.2 How to Run Simulations from Leaf Nodes: Leaf nodes can be scored with value estimates r(ˆx_0(x_t−1)), or more accurately after rolling out the pretrained model for k additional steps to evaluate r(ˆx_0(x_t−k)).The rollout brings the state closer to time 0; Figure 8 illustrates this evaluation strategy with beam search.
  • 6.2 How to Run Simulations from Leaf Nodes: Longer look-ahead improves leaf-evaluation accuracy but increases computation, with a full t-step rollout being most accurate yet computationally prohibitive.The trade-off between evaluation accuracy and computational efficiency is left for future work.

7 Editing and Refinement with Diffusion Models

Section 7 adapts inference-time diffusion techniques from generating samples from noise to editing existing designs and refining generated ones. Its iterative method combines controlled noising, reward-guided denoising, and constraint-based selection, linking the procedure to evolutionary algorithms.

  • 7 Editing and Refinement with Diffusion Models: Editing preserves existing sequence properties while enhancing target properties, whereas refinement further improves designs generated from complete noise.These use cases arise because direct from-scratch inference-time techniques do not directly support limited-mutation editing, such as antibody design with typically 5–10 mutations.
  • 7.1 Iterative Refinement in Diffusion Models: Iterative refinement repeatedly noises designs from x_0 to x_k, applies inference-time alignment back to x_0, and retains high-reward designs satisfying constraints.Constraints can include an edit-distance threshold relative to predefined seed sequences; choosing moderate k closer to 0 improves adherence when constraints are difficult, while k = T +1 recovers prior from-scratch techniques and may perform poorly.
  • 7.2 Connection with Evolutionary Algorithms: Algorithm 12 is an evolutionary algorithm variant whose inference-time mutation generates natural-like candidates more strategically and efficiently than random mutation before reward- and constraint-based selection.Its mutation stage corresponds to noising and inference-time alignment, while its selection stage retains designs meeting target rewards and constraints.

8 Comparison with Inference-Time Techniques in Language Models

The section compares inference-time guidance across diffusion, autoregressive, and masked language models, emphasizing diffusion’s forward processes, training-free value functions, and continuous-time formulations. It also explains how diffusion-style alignment can be adapted to masked language models despite decoding and distribution-shift challenges.

  • 8.1 Inference-Time Alignment Technique in Autoregressive Models: Most diffusion guidance methods transfer directly to autoregressive models, but training-free value approximations and continuous-time formulations are generally unavailable there.Analogous derivative-free and derivative-based methods have been explored for autoregressive pretrained models, as summarized in Table 1.
  • 8.2.1 Similarities and Differences Compared to Masked Diffusion Models: Masked diffusion models avoid distributional shifts more effectively than standard masked language models because their decoding process preserves matching training and inference-time distributions.Masked language models have no universally prescribed decoding strategy, so their learned encoders may mismatch the distribution induced during inference.
  • 8.1.1 Properties Leveraged in Diffusion Models: Diffusion models uniquely provide forward noising processes, training-free value-function approximations, and useful continuous-time formulations for inference-time guidance.One-step denoising from x_t to x_0 simplifies value estimation, whereas autoregressive models typically require Monte Carlo regression or soft Q-learning.
  • 8.2.1 Similarities and Differences Compared to Masked Diffusion Models: Masked diffusion models also support formally computable likelihoods through the ELBO bound, whereas masked-language-model likelihood estimates depend on decoding strategies and lack formal guarantees.Pseudo-likelihood and Monte Carlo approaches are available for masked language models, but these methods can still suffer from distributional shift.
  • 8.2.2 Adaptation of Alignment Methods: Diffusion-style alignment for masked language models defines a pretrained denoising policy by selecting masked positions and tokens using confidence, random, or left-to-right decoding.Token selection can use the encoder output with top-K sampling after choosing which position to unmask.
  • 8.2.2 Adaptation of Alignment Methods: The corresponding masked-language-model value function evaluates the single-step unmasked prediction, enabling SMC-based guidance or value-based sampling to approximate the next-step policy.For a masked state x_t, the value function is defined as r(x̂_0(x_t)), where x̂_0(x_t) is the one-step reconstruction.

9 Combining Fine-Tuning with Inference-Time Techniques

The section addresses the inference-time cost of repeatedly evaluating value functions or their derivatives by proposing policy distillation and related diffusion-distillation methods to accelerate aligned generation. It contrasts these approaches with classifier-free and RL-based fine-tuning, explaining their objectives, practical tradeoffs, and connections.

  • 9.1.1 Classifier-Free Fine-Tuning: Classifier-free fine-tuning constructs reward-labeled data and conditions a fine-tuned model on high rewards, but becomes sample-inefficient and computationally costly, especially for conditional models with multiple rewards.Inference-time techniques avoid these data-augmentation and training burdens, while also supporting classifier-free fine-tuning during data construction.
  • 9.1.2 RL-Based Fine-Tuning: Inference-time techniques are generally more stable and effective than RL-based fine-tuning for difficult non-differentiable molecular rewards because they guide individual samples without altering the underlying diffusion model.RL-based fine-tuning must optimize the entire generative model, whereas inference-time methods directly sample toward the target policy; Li et al. (2024) demonstrates effectiveness in molecular design.
  • 9.2 Policy Distillation: Policy distillation fine-tunes diffusion models to replicate inference-time teacher trajectories, iteratively improving student policies while reducing the cost of repeated guidance computations.The section also connects this approach to RL-based fine-tuning and subsequent diffusion-model methods that reduce inference steps without reward maximization.
  • 9.2.1 Choice of Roll-In Distributions: Roll-in distributions determine where policy-distillation approximation is accurate: teacher policies provide the target but can cause distribution shift, student policies mitigate that shift, and forward-process recycling is faster.Teacher and student distributions may also be mixed, while forward recycling samples from x_0 to x_t rather than the more expensive x_T to x_t.
  • 9.2.2 Choice of Divergence: KL divergence provides conservative target coverage, whereas inverse KL is more mode-seeking and benefits optimization when diversity is not the priority.The mode-seeking behavior can be disadvantageous when broad distributional coverage is required.
  • 9.3 Relation to RL-Based Fine-Tuning: KL-based policy distillation corresponds to value-weighted maximum likelihood estimation and Uehara et al. (2024, Algorithm 3), while path-consistency learning is generally more stable than inverse KL.PCL corresponds to Uehara et al. (2024, Algorithm 5) and avoids an expectation that depends on the optimized parameter.
  • 9.3.3 PPO and Direct Backpropagation: PPO and direct backpropagation target the same optimal policy as inference-time methods under ideal conditions, but inverse-KL distillation explicitly estimates value functions whereas PPO does not.Direct backpropagation is feasible when rewards are differentiable; PPO-based fine-tuning remains difficult when molecular rewards are non-differentiable.
  • 9.4 Differences between Policy Distillation and RL-Based Fine-Tuning: Policy distillation offers more stable updates and greater robustness to reward over-optimization than PPO, while standard diffusion distillation can subsequently reduce reverse steps without reward maximization.KL distillation uses fixed targets and can remain offline; its inference-time roll-ins are more likely to stay on natural manifolds, whereas PPO may drift toward out-of-distribution samples.

10 More Related Works

This section surveys additional inference-time topics, including protein-design sampling and hallucination, inpainting and inverse problems, and speculative decoding. It highlights how these methods relate to reward optimization, constrained generation, and faster inference.

  • 10 More Related Works: The section broadens the review beyond its primary focus by outlining related inference-time methods and applications not covered in earlier discussions.These topics connect reward-guided generation with protein-design refinement, constrained reconstruction, inverse problems, and inference acceleration.
  • 10.1 Walk-Jump Sampling for Protein Design: Walk-Jump Sampling targets high-reward, natural protein designs, especially for antibody design, by alternating reward-gradient “walk” steps with score-based “jump” steps.It can be viewed as single-noise-level classifier guidance whose score and gradient are iteratively updated through MCMC rather than sequentially across diffusion timesteps.
  • 10.2 Hallucination Approaches for Protein Design: Protein hallucination methods iteratively refine sequences using MCMC, evolutionary, or gradient-based algorithms to optimize structural, stability, binding, and geometric rewards.Common objectives compare predicted structures from models such as AlphaFold or ESMFold with target structures.
  • 10.3 Inference-Time Techniques for Inpainting and Linear Inverse Problems.: Inpainting methods fill missing image regions consistently, while protein motif scaffolding applies analogous conditioning; SMC-based refinement improved on heuristic replacement methods and succeeded in motif scaffolding.Specialized algorithms for linear inverse problems can outperform broader methods in constrained settings.
  • 10.4 Speculative Decoding.: Speculative decoding accelerates NLP inference by having a smaller model propose multiple tokens that a larger model validates, suggesting a similar opportunity for large-scale protein design.The proposed benefit is faster inference as protein-design models grow larger.
Loading 2501.09685v2…