PaperlayerView exact revision on arXiv

Source-linked AI summary

Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion

Jaihoon Kim, Taehoon Yoon, Prin Phunyaphibarn, Seungjun Kim, Morteza Mardani, Minhyuk Sung

arXiv:2605.23346v1cs.LG

TL;DR

Reward-tilted sampling in discrete diffusion is hindered by costly Monte Carlo estimation of the optimal twist function. CDM learns that twist contrastively, using diffusion-native forward kernels for efficient training, and achieves superior scaling across multiple applications while retaining scope limitations tied to discrete guidance and distribution mismatch.

  • Problem

    Discrete diffusion lacks the continuous-domain shortcut for twist estimation, leaving costly Monte Carlo approximation as the standard approach for reward-tilted sampling.

  • Method

    CDM learns a parameterized twist with positive and negative samples and uses closed-form diffusion forward kernels to train it efficiently.

  • Results

    CDM consistently achieves superior scaling and outperforms baselines across toxic text, regulatory DNA, protein generation, and dLLM preference alignment.

  • Takeaways & Limitations

    The learned twist amortizes SMC inference by reducing twist evaluation to a constant-time operation and can be paired with different proposal distributions.

  • Takeaways & Limitations

    Gradient-based discrete guidance remains biased and unstable and cannot handle non-differentiable objectives, while regression-based twist learning can suffer train-test distribution mismatch.

Abstract

from arXiv · show

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.

1 Introduction

Discrete diffusion enables structured generation and reward alignment, but discrete-state twist estimation makes exact reward-tilted sampling costly. CDM amortizes this inference cost through contrastive twist learning and demonstrates strong scaling across diverse applications.

  • Motivation: Reward alignment samples from a pretrained model distribution tilted toward higher human-preference or protein-designability rewards.
  • Motivation: Discrete reward alignment faces biased, unstable gradient approximations and costly Monte Carlo twist estimation, especially with expensive rewards.
  • Contrastive Distribution Matching: CDM learns the twist function contrastively, using positive and negative samples to upweight high-reward regions and downweight suboptimal ones.
  • Contrastive Distribution Matching: CDM uses the diffusion forward process and closed-form kernels to reuse clean samples across timesteps and gradient updates.
  • Empirical scope: CDM shows superior scaling and outperforms baselines across toxic text generation, regulatory DNA design, protein generation, and dLLM preference alignment.

2 Preliminary: Discrete Diffusion

Discrete diffusion models corrupt categorical data through a masking process and generate samples by simulating a learned reverse process. The forward process interpolates from data to a mask prior under a decreasing noise schedule.

  • Forward process: Masked Diffusion Models define a discrete forward corruption process that terminates in a mask state.
  • Forward process: The forward process interpolates between the data distribution p0 = pdata and the mask prior pT = m using a monotonically decreasing schedule αt.
  • Forward process: The forward transition samples each categorical state from a distribution mixing the clean one-hot vector with the mask state according to αt.
  • Reverse process: Sampling simulates the reverse process, whose exact posterior specifies transitions conditioned on the current state and, when masked, the clean state.
  • Reverse process: Because the clean data is unknown during sampling, a denoising network predicts it to parameterize the reverse transition kernel and trajectory distribution.

3 Twisted Sequential Monte Carlo for Reward Alignment

Twisted SMC samples reward-tilted discrete diffusion trajectories by combining a base model with an optimal twist, but estimating that twist is computationally difficult. CDM addresses this bottleneck by learning the twist for constant-time evaluation.

  • 3.1 KL Regularized Reward Alignment: Reward alignment maximizes expected reward while penalizing trajectory-distribution divergence from the pretrained base model through KL regularization.
  • 3.1 KL Regularized Reward Alignment: The optimal intermediate target reweights the base marginal by an exponentiated value function, yielding the optimal twist function ψ*t.
  • 3.2 Twisted Sequential Monte Carlo: Twisted SMC importance weights combine the twist ratio with the base-to-proposal transition ratio under a Markov target assumption.
  • 3.2 Twisted Sequential Monte Carlo: The optimal proposal minimizes importance-weight variance but is generally intractable because its normalizer requires evaluating the twist over all possible next states.
  • 3.2 Twisted Sequential Monte Carlo: Gradient-based discrete guidance suffers from Gumbel-Softmax bias and instability and is incompatible with non-differentiable objectives.
  • 3.2.1 Motivation: Monte Carlo Twist Function Estimation: Monte Carlo twist estimation averages exponentiated rewards from M posterior samples; increasing M improves accuracy but adds significant inference overhead.
  • 3.2.1 Motivation: Monte Carlo Twist Function Estimation: CDM trains a network to predict the optimal twist in one forward pass, reducing twist evaluation to a constant-time operation independent of proposal choice.

4 Amortized SMC with a Learned Twist Function

This section introduces CDM, which learns a twist function by matching its induced distribution to the optimal target, using contrastive positive and negative samples. A diffusion-specific forward-kernel reformulation enables efficient training and amortizes inference-time twist estimation.

  • Contrastive Distribution Matching: Regression-based twist learning trains on base-distribution samples, creating a train-test mismatch that can make targets inaccurate in regions relevant to target sampling.The mismatch motivates a distribution-level matching objective.
  • Contrastive Distribution Matching: CDM minimizes a time-averaged forward KL divergence between the optimal target distribution and the distribution induced by the parameterized twist.The induced distribution is the normalized product of the base marginal and the parameterized twist.
  • Contrastive Distribution Matching: The CDM gradient increases the twist on target-distribution samples and decreases it on samples from the current approximation, combining positive and negative learning signals.The positive term addresses distributional mismatch, while the negative term suppresses suboptimal samples.
  • Efficient Twist Parameterization: CDM amortizes costly twist estimation by learning a parameterized twist that can be evaluated with a single forward pass rather than repeated Monte Carlo computation.The learned twist is compatible with importance sampling or SMC-based negative sampling, and its inference cost is the focus of the parameterization.
  • Efficient Twist Training: Importance sampling is computationally efficient but has high variance for positive samples, whereas SMC improves alignment through resampling at the cost of sequential reward-model queries.The trade-off is especially consequential when reward evaluation is expensive.
  • Efficient Twist Training: Positive samples can be generated efficiently because the intermediate target decomposes through the closed-form diffusion forward kernel.This decomposition is described as a structural advantage of diffusion frameworks and yields an unbiased forward-based gradient estimator.

5 Related Work

Related work spans proposal fine-tuning, inference-time scaling, SMC for generative modeling, and learned twist functions. CDM targets twist-function inference cost and is positioned as complementary to proposal fine-tuning.

  • Proposal Fine-Tuning: Proposal fine-tuning methods adapt discrete diffusion alignment using specialized objectives, including mean-field GRPO in d1 and related reward-aware proposals.These methods represent a training-based alternative to directly learning the twist function.
  • Inference-Time Scaling: Inference-time scaling includes gradient guidance and search-based methods, but discrete guidance is limited by linearity assumptions and nondifferentiable rewards.The passage contrasts gradient-based guidance with search-based trajectory allocation.
  • Sequential Monte Carlo for Generative Modeling: SMC has been applied to reward-tilted sampling in continuous and discrete diffusion, while proposal fine-tuning amortizes proposal costs but leaves expensive twist estimation unresolved.CDM addresses this remaining inference bottleneck and can be combined with proposal fine-tuning.
  • Learning Twist Functions: Learned twist functions have been studied in autoregressive modeling and planning, including soft Q-learning extensions to diffusion; CDM develops a diffusion-specific contrastive objective.The contrastive formulation differs from prior work designed specifically for sequential text generation.

6 Experiments

CDM is evaluated across reward-alignment tasks using wall-clock-matched scaling, given rewards, and heldout rewards. It consistently delivers the strongest scaling, including with expensive rewards and fine-tuned proposals.

  • Toxic Text Generation: CDM achieves the best scaling on toxic text generation, outperforming all baselines for both given and heldout rewards.
  • Regulatory DNA Sequence Design: CDM maintains superior DNA-design scaling; only SMC M = 4 is comparable on the given reward, while CDM remains better on the heldout reward.
  • Protein Designability: CDM scales better than BoN and Soft Value on protein designability across given and heldout rewards despite expensive reward evaluations.
  • Diffusion LLM Alignment: CDM consistently outperforms BoN for diffusion LLM alignment, with an even larger advantage on the heldout reward.
  • Fine-Tuned Proposals: CDM paired with fine-tuned proposals outperforms the evaluated alternatives on toxic text and protein generation while mitigating mode collapse with comparable rewards.

7 Conclusion

The paper presents CDM as an amortized approach to twisted SMC for discrete diffusion, replacing costly inference-time twist estimation with a learned twist function. It reports superior scaling and training efficiency, while identifying complex rewards and extremely sparse signals as remaining challenges.

  • CDM enables amortized SMC inference in discrete diffusion by learning twist functions through a contrastive objective.
  • CDM addresses costly Monte Carlo twist estimation and achieves superior scaling performance and more efficient training than baselines.
  • CDM also provides synergistic gains when integrated with fine-tuned proposals.
  • Highly complex rewards may require more advanced twist-function architectures, while extremely sparse signals remain challenging.

A.2 Contrastive Gradient Derivations

The derivation reformulates CDM’s positive gradient term using the discrete diffusion forward process, yielding a contrastive estimator based on target and parameterized tilted distributions.

  • The contrastive gradient is the difference between expectations of ∇_ϕ log ψ_t^ϕ under the target and parameterized tilted distributions.
  • The target trajectory distribution factorizes into a reward-tilted clean marginal and the base conditional diffusion trajectory.
  • This factorization makes the target intermediate marginal equal to a clean target sample passed through the base forward kernel.
  • Consequently, the positive gradient term can be estimated by sampling clean states from the target and noisy states from the closed-form base forward process.
  • The learned twist is used in amortized inference, replacing repeated Monte Carlo twist estimation with a single forward pass.

B.2 CDM: Twist Training

CDM trains a twist network efficiently using positive and negative samples generated with the diffusion forward process, then evaluates it through a lightweight head during sampling. The method supports multiple twist-head architectures and adds less than 5% runtime overhead in reported configurations.

  • CDM Training: CDM estimates its contrastive training gradient efficiently by exploiting the diffusion forward process.The procedure is specified as contrastive twist training via the diffusion forward kernel.
  • CDM Training: Positive samples are drawn with SMC, stored in a buffer, and re-noised through the cheap forward kernel for reuse across gradient updates.The buffer is refreshed according to nupdate, while the negative term is formed separately.
  • Twist Head Parameterization: The twist function is parameterized as a lightweight head that maps the denoising network’s last hidden state to a scalar value.The design includes MLP, MLP+PE, and Transformer alternatives.
  • Twist Head Parameterization: The twist head shares the denoising backbone’s features with the logit head, allowing both outputs in one forward pass.The reported overhead is less than 5% of base-model runtime and as little as 0.5% in some configurations.
  • Experimental Setup: CDM and Soft Value use the same twist-head architecture and shared training hyperparameters for fair evaluation.Experiments cover toxic text, regulatory DNA, protein designability, and dLLM alignment, with task-specific settings.

D Ablation Studies

The ablations examine positive-buffer refresh frequency, Monte Carlo sample size, and positive-sampling method. CDM remains effective across buffer intervals and sample sizes, while SMC positive sampling outperforms importance sampling under matched settings.

  • Positive Buffer: Increasing nupdate improves training efficiency by reusing positive SMC samples and reducing reward evaluations when rewards are expensive.CDM performs well across various update intervals.
  • Monte Carlo Sample Size: CDM maintains consistent performance across different Monte Carlo sample sizes M used for target estimation during training.Soft Value instead gains initially as M increases.
  • Choice of Sampling: SMC positive sampling consistently yields better downstream performance than importance sampling under identical hyperparameter settings.The authors attribute this to SMC resampling mitigating weight degeneracy observed in importance sampling.

E Additional Results

Additional experiments test CDM with a directly fine-tuned proposal and compare its training dynamics with Soft Value. CDM improves DRAKES further and converges more efficiently to a higher final reward in regulatory DNA design.

  • Compatibility with Fine-Tuned Proposals: Applying CDM on top of DRAKES further improves performance and outperforms all other baselines.This result supports compatibility between CDM and proposals fine-tuned by direct backpropagation for differentiable rewards.
  • Training Convergence: In regulatory DNA design, Soft Value plateaus early, whereas CDM converges more efficiently and achieves a higher final reward.The comparison is presented through training dynamics rather than only final performance.

F Qualitative Results

Qualitative evaluations show closely matching generated and refolded protein structures and responses aligned with user intents across diverse dLLM tasks. The section also includes reward-training and additional qualitative comparison figures and tables.

  • Training Comparison: The section includes a training-comparison figure for regulatory DNA sequence design showing CDM and Soft Value across different M values.The supplied figure description identifies the comparison but does not state an outcome beyond the broader section evidence.
  • Protein Designability: CDM produces designable proteins whose generated and ESMFold-refolded structures closely match.The generated structure is shown in blue and the refolded structure in orange.
  • dLLM Alignment: CDM generates dLLM responses highly aligned with user intents across grammar, coding, and reasoning domains.The qualitative examples use the LLaDA-8B-Instruct base model.
Loading 2605.23346v1…