基于离散扩散模型的摊销序列蒙特卡洛对比分布匹配
阅读原文· arxiv.org离散扩散模型在生成结构化分类数据时面临从奖励倾斜分布中高效采样的挑战。扭曲序列蒙特卡洛(SMC)虽能实现渐近精确采样,但其在离散状态空间中估计最优扭曲函数需要昂贵的蒙特卡洛近似,成为推理瓶颈。为此,本文提出对比分布匹配(CDM)框架,通过学习一个参数化扭曲函数来摊销SMC推理的成本。训练时,梯度估计器被重新设计以利用离散扩散模型的闭式前向核。实验表明,评估该扭曲函数带来的额外计算开销低于基础模型单次前向传播的5%。在匹配实际耗时的条件下,CDM性能优于现有基线,并在毒性文本生成、调控DNA序列设计、蛋白质可设计性及扩散大语言模型对齐等多个任务中验证了其有效性。
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.