VASR: Variance-Aware Systematic Resampling for Reward-Guided Diffusion

TL;DR

VASR introduces a variance-aware resampling method for reward-guided diffusion models, significantly improving sample diversity and quality while maintaining computational efficiency.

Contribution

The paper proposes VASR, a novel variance decomposition framework and resampling technique that enhances reward-guided diffusion SMC performance without additional training.

Findings

VASR achieves up to 26% better FID on MNIST and CIFAR-10.

VASR is 66 times faster than MCTS-based methods at similar compute levels.

VASR-Max outperforms prior SMC baselines in text-to-image generation.

Abstract

Sequential Monte Carlo (SMC) samplers for reward-guided diffusion models often suffer from rapid lineage collapse: a few high-reward particles dominate the population within a handful of resampling steps, destroying diversity and degrading sample quality. We propose a variance-decomposition framework for reward-guided diffusion SMC that separates continuation variance $V_t^{\mathrm{cont}}$ from residual variance $V_t^{\mathrm{res}}$, revealing that high offspring-count variance under the commonly used multinomial resampling drives this collapse. This motivates \textsc{VASR} (Variance-Aware Systematic Resampling), which addresses both variance terms via variance-optimal mass allocation $m_t \propto w_t e^{r_t}$ (minimizing $V_t^{\mathrm{cont}}$) and systematic resampling (controlling $V_t^{\mathrm{res}}$). For latent diffusion models where intermediate rewards are noisy due to stochastic continuations, we propose \textsc{VASR-Max}, a deliberately biased high-selection variant for variance-sensitive reward optimization. Both methods are training-free, fully parallelizable, and add only linear overhead. On MNIST and CIFAR-10, VASR achieves as high as $26\%$ better FID than prior SMC methods while remaining 66 times faster than MCTS-based value methods at matched compute. On text-to-image generation, \textsc{VASR-Max} consistently outperforms the strongest SMC baseline across compute budgets and matches MCTS-based methods within 2.5--3% reward at high budgets while being approximately times faster.