Discrete Diffusion Inference-Time Control With Nested Sequential Monte Carlo

Lohithsai Yadala Chanchu, Hany Abdulsamad, Christian A. Naesseth

Keywords: Discrete Diffusion Language Models, Inference-time Control, Feynman-Kac Steering, Reward-Guided Generation, Diffusion Models, Steering, Fine-tuning, Particle-based Sampling

TL;DR: The authors propose Nested SMC and FA-NSMC for steering discrete diffusion LMs. Correcting prior biases , they outperform Best-of-N and bootstrap SMC on toxicity and fluency tasks.

Abstract: We study inference-time control for text generation in discrete diffusion language models, where the goal is to steer sampling toward sequence-level rewards without retraining. Prior work in this domain has focused on particle-based methods such as best-of-$n$ sampling and bootstrap sequential Monte Carlo, which may suffer from overoptimism and weight degeneracy, respectively. We address these limitations using \emph{nested} sequential Monte Carlo methods. We formulate nested SMC (NSMC) and fully-adapted nested SMC (FA-NSMC) for Feynman--Kac steering, identifying and correcting errors in prior formulations that lead to biased final estimates. We evaluate these methods on toxicity and fluency steering tasks, showing that NSMC and FA-NSMC consistently outperform best-of-$n$ and bootstrap SMC.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 95