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

TL;DR

This work surveys inference-time alignment methods for diffusion models, focusing on how to maximize downstream rewards (e.g., stability, affinity) without fine-tuning. It presents a unified framework where soft value functions guide a pre-trained denoising process to approximate the target reward-augmented distribution, and contrasts derivative-free approaches (SMC, beam search) with derivative-based methods (classifier guidance, Doob transform). The tutorial integrates discrete and continuous diffusion formalisms, extends guidance to Riemannian manifolds for protein structures, and connects inference-time strategies to post-training and RL-based tuning, offering practical algorithms, tree-search concepts, and editing/refinement techniques. The material provides concrete pathways to apply inference-time alignment in protein design and related domains, with links to language-model analogies and forthcoming computational tools (e.g., AlignInversePro).

Abstract

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

Figures (9)

• Figure 1: The objective of inference-time techniques is to generate natural designs (e.g., natural images or natural-like protein sequences) with high functionality, without any direct fine-tuning of diffusion models.
• Figure 2: Summary of representative inference-time algorithms. Here, we aim to optimize downstream reward functions $r:\mathcal{X} \to \mathbb{R}$ given pre-trained masked diffusion models for sequences. The value function $v(\cdot)$ serves as a look-ahead function, mapping intermediate states to expected future rewards $r(\cdot)$. Best-of-N is a naïve method that selects the best sample among $N$ generated ones. Derivative-based guidance adds gradients of differentiable value function models during inference, making it a powerful method when we can construct the actual value function models. SMC-Based Guidance and Value-Based Importance Sampling (a.k.a. beam search with value functions) are gradient-free methods that sequentially select favorable intermediate states based on value functions. These methods do not require constructing differentiable value function models, which can often be challenging in molecular design.
• Figure 3: Scaling inference time compute via value-based beam search li2024derivative in sec:beam, progressively increasing the tree width (see fig:whoe_summaryd). These figures demonstrate that as the computational budget (x-axis) increases, rewards (y-axis) can be optimized more effectively.
• Figure 4: Roadmap of This Paper.
• Figure 5: Typical foundational diffusion models for protein Sequences (especially, Red indicates diffusion models). The former approach relies solely on sequence data, while the latter explicitly generates the structure first, followed by the sequence. Notably, hybrid approaches that combine these methods have become increasingly popular in recent studies.

Theorems & Definitions (12)

• Example 1: Euclidean space
• Remark 1: Different parameterization
• Example 2: Discrete space (masked diffusion models)
• Remark 2
• Theorem 1: From Theorem 1 in uehara2024bridging
• Theorem 2: Doob Transform
• Remark 3
• Example 3: SO(3)
• Example 3: continued
• Remark 4: Combination with More Advanced Discretization Methods