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

TL;DR

This paper introduces LiDAR sampling, a test-time scaling method for diffusion models that efficiently guides sample generation towards human-aligned rewards using a novel closed-form guidance computation and lookahead sampling, achieving high performance with less computation.

Contribution

It proposes a new EFR formulation that enables closed-form guidance without neural backpropagation and introduces LiDAR sampling with lookahead to improve efficiency and reward alignment.

Findings

LiDAR achieves similar performance to gradient guidance methods with fewer samples.

Steep performance gains with increased lookahead accuracy and sample count.

LiDAR reaches SDXL GenEval performance with 3 samples and 3-step lookahead, 9.5x faster.

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.