arXiv:2605.23346cs.LG2026-05

用对比学习加速离散扩散模型的采样,提升生成效率与质量。

Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion

论文配图:Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion
图 1 · 摘自论文原文
  • 通过正负样本学习参数化扭曲函数,降低采样成本。
  • 推理时额外开销低于5%,性能超越现有方法。
  • 适用于文本、基因序列、蛋白质等结构化数据生成。

离散扩散模型在生成结构化类别数据方面表现出强大能力。然而,高效采样奖励倾斜分布仍是核心挑战。尽管扭曲序列蒙特卡洛(Twisted SMC)可实现渐近精确,但在离散状态空间中估计最优扭曲函数需依赖昂贵的蒙特卡洛近似,造成严重计算瓶颈。为此,我们提出对比分布匹配(CDM),通过正负样本学习参数化扭曲函数,实现SMC推断的费用摊销。为高效训练,我们重构梯度估计器,利用离散扩散模型的闭式前向核。实际应用中,评估所学扭曲函数的额外开销不足基模型单次前向传播的5%。大量实证表明,CDM在相同运行时间下持续优于现有基线。我们在有毒文本生成、合规DNA序列设计、蛋白质可设计性及扩散大语言模型对齐等多个场景验证了该方法的有效性与通用性。

原文摘要 · Abstract (English)

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.

离散扩散序列生成采样优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。