Lookahead Sample Reward Guidance for Test-Time Scaling of Diffusion Models

TL;DR

This paper tackles the misalignment of diffusion-model samples with human intent and the computational burden of gradient-based reward guidance. It introduces LiDAR, a lookahead-based sampling framework that derives the Expected Future Reward from future marginal samples, enabling closed-form, derivative-free guidance that does not backprop through reward models. The authors prove an error bound for lookahead approximations and demonstrate strong scaling: LiDAR achieves state-of-the-art GenEval performance on SDXL with up to ~9.5x faster inference compared to gradient-guided methods, and performs well with modest lookahead budgets (e.g., three samples). Moreover, LiDAR integrates orthogonally with SMC and Best-of-N strategies and remains robust as reward strength increases, illustrating practical impact for reward-aligned diffusion sampling.

Abstract

Diffusion models have demonstrated strong generative performance; however, generated samples often fail to fully align with human intent. This paper studies a test-time scaling method that enables sampling from regions with higher human-aligned reward values. Existing gradient guidance methods approximate the expected future reward (EFR) at an intermediate particle $\mathbf{x}_t$ using a Taylor approximation, but this approximation at each time step incurs high computational cost due to sequential neural backpropagation. We show that the EFR at any $\mathbf{x}_t$ can be computed using only marginal samples from a pre-trained diffusion model. The proposed EFR formulation detaches the neural dependency between $\mathbf{x}_t$ and the EFR, enabling closed-form guidance computation without neural backpropagation. To further improve efficiency, we introduce lookahead sampling to collect marginal samples. For final sample generation, we use an accurate solver that guides particles toward high-reward lookahead samples. We refer to this sampling scheme as LiDAR sampling. LiDAR achieves substantial performance improvements using only three samples with a 3-step lookahead solver, exhibiting steep performance gains as lookahead accuracy and sample count increase; notably, it reaches the same GenEval performance as the latest gradient guidance method for SDXL with a 9.5x speedup.

Figures (15)

• Figure 1: Overview of the proposed sampling framework from prompt ${\mathbf{c}}$. (a) To obtain marginal samples, we generate lookahead samples using a few-step solver and annotate them with a reward function. (b) To generate samples from the target reward-tilted distribution, LiDAR guides particles ${\mathbf{x}}_t$ toward high-reward lookahead sample $\hat{{\mathbf{x}}}_0^{1}$ and repels low-reward lookahead samples $\hat{{\mathbf{x}}}_0^{2}, \hat{{\mathbf{x}}}_0^{3}$. The guiding weight $w_i^r - w_i$ provided by each lookahead sample $\hat{{\mathbf{x}}}_0^{i}$ is proportional to its reward annotation value; see \ref{['eq:empirical_guidance']} for the definition.
• Figure 2: Sampling results for the prompt "a photo of a yellow bird and a black motorcycle." Among the lookahead samples, the blue boxes indicate all samples used for LiDAR (DPM-3, $n$=3), ordered from left to right by increasing reward values. The green box indicates the lookahead sample with the highest reward used for LiDAR (DPM-5, $n$=50).
• Figure 3: Scaling behavior of the LiDAR sampler under different lookahead strategies and varying numbers of lookahead samples $n$. The vanilla method uses only the Stein score ${\mathbf{s}_{\boldsymbol{\theta}}}$, whereas LiDAR (w/ DPM-$\delta$) incorporates lookahead samples with $\delta$ discretization steps. All results are obtained using SD v1.5 with ImageReward annotations. The legend is shared across all three panels.
• Figure 4: (a) Performance trade-offs between UG and LiDAR. (b, c) Efficiency–performance trade-offs for test-time scaling methods. Vanilla increases sampling steps, DATE adjusts gradient update frequency, and LiDAR controls $n$. The cross marker indicates the performance in \ref{['tab:main2']}. The black dotted line indicates the time difference for LiDAR to reach DATE’s maximum scaling performance.
• Figure 5: Performance comparison with other sample-based guidance methods using our lookahead samples based on DPM-5 lookahead solver with $n=50$.

Theorems & Definitions (10)

• Theorem 3.1
• Definition 3.1
• Theorem 3.2
• Theorem 3.3
• Theorem 1.1
• proof
• Theorem 1.1
• proof
• Theorem 1.1
• proof