让扩散语言模型学会像推理时那样逐步去噪,提升生成质量。
Self-Distilled Trajectory-Aware Boltzmann Modeling: Bridging the Training-Inference Discrepancy in Diffusion Language Models

- 用玻尔兹曼分布建模推理过程中的去噪偏好,对齐训练与推理路径。
- 在新领域任务上显著提升性能,缓解灾难性遗忘问题。
- 适合追求高质量生成、重视推理一致性的扩散模型研究者。
扩散语言模型(DLMs)作为自回归模型的替代方案,具备更强的全局感知和高度并行生成能力。然而,基于标准负证据下界(NELBO)的监督微调在后训练阶段效率低下:训练时单步重构随机掩码的词元,而推理时遵循信心引导的多步从易到难去噪轨迹。现有基于轨迹的自蒸馏方法主要用以压缩采样步骤、加速推理,常无法实质性提升模型能力,甚至在完整扩散解码下导致性能下降。本文提出轨迹对齐的玻尔兹曼建模(TABOM),旨在利用自蒸馏轨迹实现真正意义上的知识获取。尽管这些轨迹位于预训练DLM自身的分布流形上,优化门槛较低,但直接使用标准NELBO微调仅带来边际提升。TABOM将推理去噪偏好建模为预测熵上的玻尔兹曼分布,并推导出可计算的成对排序目标,使模型的置信度排序与观测到的解码轨迹对齐。实验证明,TABOM在新领域中取得显著提升,扩展了DLM的有效知识边界,并显著缓解了灾难性遗忘问题。
原文摘要 · Abstract (English)
Diffusion Language Models (DLMs) have recently emerged as a promising alternative to autoregressive language models, offering stronger global awareness and highly parallel generation. However, post-training DLMs with standard Negative Evidence Lower Bound (NELBO)-based supervised fine-tuning remains inefficient: training reconstructs randomly masked tokens in a single step, whereas inference follows a confidence-guided, multi-step easy-to-hard denoising trajectory. Recent trajectory-based self-distillation methods exploit such inference trajectories mainly for sampling-step compression and acceleration, often improving decoding efficiency without substantially enhancing the model's underlying capability, and may even degrade performance under full diffusion decoding. In this work, we ask whether self-distilled trajectories can be used not merely for faster inference, but for genuine knowledge acquisition. Although these trajectories lie on the pretrained DLM's own distributional manifold and thus offer a potentially lower optimization barrier, we find that naively fine-tuning on them with standard NELBO objectives yields only marginal gains. To address this limitation, we propose \textbf{T}rajectory-\textbf{A}ligned optimization via \textbf{Bo}ltzmann \textbf{M}odeling (\textbf{TABOM}), a self-distilled trajectory-based post-training framework that aligns training with the easy-to-hard structure of inference. TABOM models the inference unmasking preference as a Boltzmann distribution over predictive entropies and derives a tractable pairwise ranking objective to align the model's certainty ordering with the observed decoding trajectory. Empirically, TABOM achieves substantial gains in new domains, expands the effective knowledge boundary of DLMs, and significantly mitigates catastrophic forgetting compared with standard SFT.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。