Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion
For practitioners using discrete diffusion models, this method provides a practical and efficient way to perform reward-guided sampling without the heavy computational cost of traditional SMC.
The paper tackles the computational bottleneck of Twisted Sequential Monte Carlo (SMC) for sampling from reward-tilted distributions in discrete diffusion models. The proposed Contrastive Distribution Matching (CDM) learns a parameterized twist function with less than 5% overhead, consistently outperforming baselines under matched wall-clock time across tasks like toxic text generation and DNA sequence design.
Discrete diffusion models have emerged as powerful frameworks for generating structured categorical data. However, efficiently sampling from reward-tilted distributions remains a fundamental challenge. While Twisted Sequential Monte Carlo (SMC) offers asymptotic exactness for this task, estimating the optimal twist function in discrete state spaces necessitates costly Monte Carlo approximations, resulting a severe computational bottleneck at inference. To overcome this limitation, we introduce Contrastive Distribution Matching (CDM), a novel framework that amortizes the cost of SMC inference by learning a parameterized twist function via positive and negative samples. For efficient training, we reformulate the gradient estimator to leverage the closed-form forward kernels of discrete diffusion models. In practice, evaluating our learned twist function incurs less than 5% additional computational overhead compared to a single forward pass of the base model. Through extensive empirical evaluations, we demonstrate that CDM consistently outperforms existing baselines under matched wall-clock time. We validate the effectiveness and versatility of our approach across a diverse range of applications, including toxic text generation, regulatory DNA sequence design, protein designability, and diffusion large language model alignment.