Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion

Abstract: 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 LLM alignment.

• The paper presents CDM, which uses a contrastive objective to learn a twist function and replace costly Monte Carlo estimates in reward-biased sampling.
• It leverages closed-form diffusion kernels and efficient positive/negative sampling to achieve constant-time inference with a lightweight neural head.
• Empirical results across text, DNA, and protein design demonstrate CDM’s superior convergence, scaling rewards, and reduced computational overhead.

Contrastive Distribution Matching for Amortized SMC in Discrete Diffusion

Introduction and Motivation

The paper "Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion" (2605.23346) introduces a principled framework for reward-tilted sampling in discrete diffusion models via amortized Sequential Monte Carlo (SMC). Discrete diffusion models excel at generating categorical and structured data across language, DNA, and protein domains, yet inference for reward-biased distributions is computationally prohibitive due to the requirement of Monte Carlo estimates for the optimal twist function. The authors present Contrastive Distribution Matching (CDM), which learns a parameterized twist function using a contrastive objective, enabling constant-time, forward-pass inference and decomposing training into efficient positive/negative sampling via closed-form diffusion kernels.

Discrete Diffusion and Reward Alignment

Discrete diffusion models parameterize a forward masking process with categorical transitions, interpolating data and a terminal mask through a noise schedule. Sampling proceeds by reverse diffusion, using a denoising network to approximate posterior transitions. Reward alignment reweights the diffusion prior toward high-reward regions, formulating the optimal target distribution with twist functions derived from value functions over reward expectations.

Twisted SMC enables asymptotically-correct sampling from the reward-tilted target by sequential importance sampling and resampling, but efficient estimation of the twist function is a key challenge. Continuous diffusion benefits from Tweedie's formula for twist approximation, but discrete models lack an equivalent, needing expensive Monte Carlo sampling for the twist and incurring a severe computational bottleneck for high-cost rewards.

CDM: Amortized Twist via Contrastive Learning

Traditional regression-based approaches learn the twist by minimizing MSE between predictions and Monte Carlo twist estimates sampled from the base process. However, this leads to train/test mismatch, as the base proposal does not concentrate the training on high-reward regions. CDM circumvents this via a distribution-matching objective: minimizing the forward KL divergence between the twist-inducing distribution and the optimal target. The resulting gradient has explicit contrastive structure, upweighting positive samples from the reward-tilted target and downweighting negatives from the current model.

The authors exploit the closed-form forward diffusion process, enabling high-efficiency positive sampling: a buffer of clean samples can be renoised across timesteps, amortizing reward computation. This forward-based gradient estimator provides unbiased updates and overcomes the sample inefficiency of standard SMC. Negative samples are collected via importance sampling, and an EMA update stabilizes training.

Efficient Parameterization and Inference

CDM is realized as a lightweight neural head atop the diffusion backbone, outputting twist values per sequence or token. The architecture considers MLP, MLP + positional encoding (PE), and Transformer variants. Runtime analysis demonstrates that the twist head incurs less than $5\%$ of base model inference time, enabling SMC sampling at scale, compared to previous approaches requiring independent value networks.

Figure 1: (Left) CDM attaches a lightweight head to the last hidden state for efficient twist estimation. (Right) Three architectural variants: MLP, MLP+PE, and Transformer.

Empirical Evaluation and Scaling Results

Extensive experiments span text toxicity generation, regulatory DNA sequence design, protein designability, and diffusion LLM alignment. Scaling performance is assessed against wall-clock time for reward and heldout reward metrics, with CDM consistently establishing a new Pareto front in all domains.

Figure 2: CDM outperforms all baselines in scaling reward generation across four diverse tasks, setting a new Pareto front in wall-clock time.

Key observations include:

• SMC performance is strongly constrained by per-step Monte Carlo computation for twist evaluation, especially as reward complexity grows.
• Regression-based Soft Value methods amortize some costs but are suboptimal in high-reward regions due to train/test mismatch.
• CDM yields faster convergence, higher final rewards, and superior generalization to heldout metrics.
• In protein design tasks, CDM enables practical generation of designable proteins despite the expensive folding reward.

Figure 3: CDM-generated proteins (blue) are accurately refolded by ESMFold (orange), validating protein designability.

Synergy with Proposal Fine-Tuning and Mode Collapse

CDM can be paired with fine-tuned proposals, such as d1 and DRAKES. When applied atop models tuned with direct backpropagation (DRAKES) or mean-field RL objectives (d1), CDM achieves synergistic gains in reward scaling, while maintaining diversity and robust generative perplexity.

Figure 4: CDM further enhances reward scaling in conjunction with fine-tuned proposals, outperforming d1 and DRAKES.

CDM mitigates mode collapse frequently observed in RL-based fine-tuning, producing higher diversity, better self-BLEU, and higher cluster counts in protein design.

Training Dynamics and Ablation

CDM exhibits superior training convergence relative to regression-based approaches, as shown in wall-clock reward plots and training ablations. Positive buffer-based training intervals and Monte Carlo sample sizes are analyzed; CDM maintains robust performance across a range of settings.

Figure 5: CDM convergence (left) outpaces Soft Value baselines for regulatory DNA design. CDM + DRAKES (right) demonstrates compatibility and scaling gains.

Amortized SMC Inference Procedure

CDM can be plugged directly into the SMC sampling routine, replacing costly Monte Carlo twist estimation with a single neural forward pass, thus enabling efficient reward-tilted sampling at scale.

Figure 6: CDM enables amortized SMC inference with a single forward-pass twist function, in contrast to classical SMC requiring expensive Monte Carlo estimates.

Practical and Theoretical Implications

CDM provides a scalable, modular solution for reward-tilted sampling in discrete diffusion models. Practically, it enables real-world applications such as DNA regulatory sequence design and protein engineering by reducing inference cost, supports computational safety benchmarking in text generation, and aligns diffusion LLMs via API-based non-differentiable rewards. Theoretically, CDM bridges SMC inference with representation learning via contrastive distribution matching, leveraging structural advantages in diffusion forward processes unavailable in sequential autoregressive models.

Methodological advances include orthogonality with proposal fine-tuning—CDM can combine efficiently with any reward-aware proposal method—and amortization of inference cost for high-complexity objectives. These design principles are applicable to scientific generative tasks and scalable LLM alignment, suggesting a strong foundation for future hybrid training/inference approaches and generalized reward-based generative modeling.

Conclusion

Contrastive Distribution Matching (CDM) is a coherent solution to the challenge of scalable reward-biased sampling in discrete diffusion. By amortizing SMC inference through a contrastive objective and exploiting structural kernel decompositions, CDM achieves superior empirical performance and efficiency, enables modular synergy with fine-tuned proposals, and overcomes the computational bottlenecks inherent in naive Monte Carlo inference. The framework provides both practical impact for scientific and language domains and theoretical advancement in probabilistic generative modeling, setting the stage for further exploration in amortized sampling and distribution matching for discrete generative models.