解决扩散语言模型训练与推理不匹配问题,让模型适应任意文本结构。
Adaptive Block Diffusion: Resolving Training-Inference Mismatch in Diffusion Language Models

- 将文本配置设为随机变量,统一训练所有可能的上下文结构。
- 在多种解码尺度下保持性能稳定,避免离格崩溃现象。
- 无需修改架构即可适配不同推理策略,适合实际部署场景。
扩散语言模型(DLMs)通常在固定上下文结构下训练,仅对预定义的词元子集进行去噪,导致训练与推理之间存在不匹配:模型在推理时需处理任意配置,造成性能下降。本文提出自适应块扩散(Adaptive Block Diffusion, ABD),通过在前缀-窗口配置分布上优化去噪风险,解决该问题。将配置视为随机变量,使单一模型在全配置空间中训练,无需架构改动。我们证明,跨解码策略的泛化能力取决于训练分布的支持范围,且ABD能保证任何在训练中覆盖过的推理策略达到去噪最优。实验表明,ABD在不同解码尺度下具有结构不变性,避免了离格崩溃,并恢复了块大小与困惑度之间的单调关系,其性能在目标尺度上可媲美甚至超过专用固定块模型。
原文摘要 · Abstract (English)
Diffusion Language Models (DLMs) are typically trained under fixed context structures, restricting denoising to predetermined token subsets. This creates a mismatch between training and inference, where models must operate over arbitrary configurations, leading to degradation off the training grid. We propose Adaptive Block Diffusion (ABD), which resolves this mismatch by optimizing denoising risk over a distribution of prefix-window configurations. By treating the configuration as a stochastic variable, ABD trains a single model over the full configuration space without architectural changes. We show that generalization across decoding strategies is governed by the support of the training distribution, and that ABD guarantees denoising optimality for any inference policy whose configurations are covered during training. Empirically, ABD exhibits structural invariance across decoding scales, avoiding off-grid collapse and recovering a monotonic relationship between block size and perplexity, while matching or outperforming fixed-block specialists at their target scales.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。