Psi-Sampler: Initial Particle Sampling for SMC-Based Inference-Time Reward Alignment in Score Models
Taehoon Yoon, Yunhong Min, Kyeongmin Yeo, Minhyuk Sung
TL;DR
PSI-Sampler addresses the inefficiency of inference-time reward alignment in score-based generative models by initializing particles from the reward-aware posterior $\tilde{p}_1^*(\boldsymbol{x}_1)$ rather than the Gaussian prior $p_1$. It employs the preconditioned Crank–Nicolson Langevin (pCNL) sampler to draw samples from this posterior and then feeds them into Sequential Monte Carlo with the approximately optimal transition kernel $\tilde{p}_{\theta}^*$ to target $p_0^*$, using Tweedie’s formula within a stochastic optimal control framework. Through experiments on layout-to-image, quantity-aware generation, and aesthetic-preference generation, it demonstrates consistent improvements in both seen and held-out rewards over Gaussian-prior baselines and other posterior initializations. The results show that reward-informed initialization yields better exploration and sample quality under fixed compute budgets, with practical implications for reward-aligned generation in high-dimensional score-based models.
Abstract
We introduce $Ψ$-Sampler, an SMC-based framework incorporating pCNL-based initial particle sampling for effective inference-time reward alignment with a score-based generative model. Inference-time reward alignment with score-based generative models has recently gained significant traction, following a broader paradigm shift from pre-training to post-training optimization. At the core of this trend is the application of Sequential Monte Carlo (SMC) to the denoising process. However, existing methods typically initialize particles from the Gaussian prior, which inadequately captures reward-relevant regions and results in reduced sampling efficiency. We demonstrate that initializing from the reward-aware posterior significantly improves alignment performance. To enable posterior sampling in high-dimensional latent spaces, we introduce the preconditioned Crank-Nicolson Langevin (pCNL) algorithm, which combines dimension-robust proposals with gradient-informed dynamics. This approach enables efficient and scalable posterior sampling and consistently improves performance across various reward alignment tasks, including layout-to-image generation, quantity-aware generation, and aesthetic-preference generation, as demonstrated in our experiments. Project Webpage: https://psi-sampler.github.io/