🤖 machine learning

Adaptive Block Diffusion: Resolving Training-Inference Mismatch in Diffusion Language Models

本文引入了自适应块扩散（Adaptive Block Diffusion, ABD），这是一种通过在多种前缀窗口配置的分布上优化去噪，从而解决扩散语言模型中训练与推理不匹配问题的方法，进而使单个模型能够在无需改变架构的情况下，稳健地泛化至各种任意解码策略。

原作者： Gagan Jain

发布于 2026-06-30

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Gagan Jain

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

以下是关于“自适应块扩散”（Adaptive Block Diffusion）论文的解释，使用了简单的语言和日常类比。

问题所在：“僵化的蓝图” vs. “现实世界”

想象你正在教一个机器人写故事。

自回归模型（Autoregressive models）（旧方法）就像一个学生，一次只写一个词，必须等前一个词写完才能写下一个。这种方式非常准确，但速度很慢。
扩散模型（Diffusion models）（新方法）则像是一个学生，他从满篇乱码的页面开始，通过慢慢擦除错误的部分来展现故事。这种方式快得多，因为他们可以同时修正许多个词。

具体问题：
近期的“块扩散”（Block Diffusion）模型试图结合两者的优点。它们决定以块（即词组块）为单位来修正故事。例如，它们可能会说：“我们先只修正第 1–10 个词，然后是 11–20 个，接着是 21–30 个。”

陷阱：
研究人员发现，这些模型仅针对那种特定的、僵化的模式（如 1–10，11–20）进行训练。

类比： 想象你只在一条车道间距完美、笔直且空旷的高速公路上练习开车。你成了这条特定赛道的专家。但到了考试那天，你却被要求在一条充满弯道、坑洼不平且车道宽度不断变化的山路上驾驶。因为你只练习过“笔直高速公路”的模式，结果你撞车了。
在论文中： 当这些模型尝试使用不同的块大小（例如一次修正 5 个词而不是 10 个词）来生成文本时，它们的性能会崩溃。它们无法处理那些未曾练习过的“非网格化”情况。

解决方案：“自适应块扩散”（ABD）

作者提出了一种名为**自适应块扩散（Adaptive Block Diffusion, ABD）**的新训练方法。

类比：
与其教机器人只在笔直的高速公路上开车，不如带它去一所会抛出随机场景的驾驶学校。

有时，它练习一次修正 1 个词。
有时，它练习一次修正 10 个词。
有时，它练习一次修正 50 个词。
有时，“块”的起始位置也会发生变化。

通过在所有可能的块大小的随机组合上进行训练，机器人学习的是修正文本的通用技能，而不是死记硬背某种特定模式。

它是如何工作的（机制）

随机变量： 论文将“块大小”（即我们一次修正多少个词）视为一个随机变量。在训练期间，模型不知道下一个会得到多大的块。它必须随时做好应对任何情况的准备。
无需新硬件： 你不需要制造一个新的机器人。你只需要改变课程内容（训练数据分布）。模型的架构保持不变；它只是学会了变得更加灵活。
保证： 论文从数学上证明了，如果你在足够多样化的块大小上训练模型，只要该大小包含在训练组合中，它在以后使用任何块大小时都能完美运行。

结果：为什么这很重要

论文在两个主要的语言数据集（LM1B 和 OpenWebText）上进行了测试，并发现：

不再崩溃： 与那些在测试不同块大小时会失败的旧有“固定块”模型不同，ABD 模型在所有尺寸下都能平稳运行。
“单调”关系： 在理想情况下，如果你减小块的大小（即一次修正更少的词），质量应该会略微提升（趋向于那种缓慢但完美的自回归风格）。旧模型打破了这一规则；当改变大小时，它们的表现反而变差了。ABD 则完美遵循了这一规则：块越小 = 质量越高，块越大 = 速度越快。
一模通用： 你不需要为一个“快速模式”训练一个模型，再为另一个“高质量模式”训练另一个模型。一个单一的 ABD 模型就可以兼顾两者，并且在各自的任务中都能达到专门化模型的水平。

简而言之

旧方法： 训练一个只能针对特定块大小工作的模型。它是一个专家，但一旦规则改变就会失败。
新方法 (ABD)： 在所有块大小的随机混合中训练模型。它变成了一个通才，能够处理你抛给它的任何情况，且不会损失质量。

该论文声称，这解决了模型训练方式与实际使用方式之间的不匹配问题，使扩散语言模型在无需复杂架构改动的情况下，变得更加鲁棒且多功能。

技术摘要：自适应块扩散 (Adaptive Block Diffusion, ABD)

1. 问题陈述

扩散语言模型 (DLMs) 为自回归 (AR) 模型提供了一种极具前景的替代方案，通过实现并行解码，有望降低推理延迟。然而，最近的结构化扩散方法，例如掩码离散扩散语言模型 (MDLM) 和块扩散（如 BD3LM），依赖于固定的上下文结构。在这些模型中，待去噪的 token 集遵循预定的模式（例如固定的块大小和对齐的前缀）贯穿整个训练过程。

这种设计造成了根本性的训练-推理不匹配：

训练阶段： 模型针对一个狭窄且固定的配置子集（“训练网格”）优化去噪风险。
推理阶段： 实际的解码策略通常需要任意长度的前缀和窗口大小（例如滑动窗口、自适应块大小或流式传输）。

当固定块模型在超出其训练支持范围（“离网”）的配置上进行评估时，会出现困惑度（perplexity）剧烈下降和生成行为不稳定的现象。它们未能学习到通用的去噪程序，而是仅专门针对训练期间看到的特定配置进行了特化。本文指出，单一的扩散模型应当能够消除这种不匹配，从而对解码配置保持不变性。

2. 方法论：自适应块扩散 (ABD)

作者提出了自适应块扩散 (ABD)，这是一个通过将去噪配置视为随机变量而非固定超参数来解决训练-推理不匹配的框架。

核心机制

ABD 不再将训练限制在固定的块对齐子集内，而是定义了一个配置空间 $\Omega := \{(k, \ell) \mid 0 \le k \le L-1, 1 \le \ell \le L-k\}$ ，其中 $k$ 是前缀长度， $\ell$ 是去噪窗口长度。

随机训练： 模型在定义在 $\Omega$ 上的分布 $\pi$ 上进行训练。在每个训练步骤中，从 $\pi$ 中采样一个配置 $(k, \ell)$ 。
去噪过程： 由窗口 $S(k, \ell) = \{k+1, \dots, k+\ell\}$ 索引的 token 通过前向扩散过程被破坏，而前缀 $x_{1:k}$ 保持干净并作为条件上下文。
目标函数： 训练目标是最小化在分布 $\pi$ 下的期望去噪风险：
$\mathcal{L}_{ABD}(\theta) = \mathbb{E}_{(k,\ell)\sim\pi, x_0\sim D, t} \left[ \sum_{i \in S(k,\ell)} \text{KL}\left(q(x^i_0 | x^i_t, t) \parallel p_\theta(x^i_0 | x^i_t, x^i_{1:k}, t)\right) \right]$

实现细节

架构： ABD 对底层的扩散模型（如 Transformer）无需任何架构改动。它仅修改了训练过程中遇到的配置分布。
注意力掩码： 每个步骤都会根据采样的 $(k, \ell)$ 构建动态注意力掩码。该掩码确保噪声 token 仅能注意到同一块内的其他噪声 token 以及前序块中的干净 token，同时干净 token 对之前的干净 token 进行因果注意力计算。
分布选择： 本文探讨了多种 $\pi$ 的分布选择（分类指数分布、均匀分布、对数正态分布）。在主要实验中，使用基于 2 的幂次 $\{1, 2, 4, \dots, 128\}$ 的分类分布并带有指数衰减，这优先考虑了较小的块大小。

3. 理论分析

论文提供了将训练分布的支持集与泛化性能联系起来的正式保证。

风险最小化： 作者将扩散训练形式化为配置空间 $\Omega$ 上的风险最小化。
定理 1 (最优性)： 任何 ABD 目标的全局极小值点都能在 $\pi$ -几乎处处 的配置 $(k, \ell)$ 下实现最优去噪性能。对于 $\pi$ 支持集之外的配置，则保持不受约束。
定理 2 (训练-推理对齐)： 如果推理策略诱导出一个分布 $\Pi_{inf}$ ，使得 $\Pi_{inf} \ll \pi$ （即推理支持集包含在训练支持集内），那么在 $\pi$ 下训练以最小化风险的模型也是推理风险下的全局极小值点。
含义： 泛化能力不是模型架构的属性，而是训练分布的支持覆盖度的属性。固定块模型之所以失败，是因为对于任意解码策略， $\Pi_{inf} \not\ll \pi_{fixed}$ 。

4. 实验结果

作者在 LM1B (One Billion Word) 和 OpenWebText (OWT) 基准测试上评估了 ABD，并将其与自回归 (AR)、MDLM、SEDD 以及固定块专家模型 (BD3LM) 进行了对比。

关键发现

结构不变性：
- 固定块专家模型（例如在块大小为 4 时训练的 BD3LM）在评估大于其训练配置的块大小时，会出现剧烈的困惑度激增。
- ABD 在全范围内恢复了块大小与困惑度之间的单调关系。随着块大小减小（趋近于自回归精度），困惑度单调改善，且不会出现离网崩溃。
性能对等：
- 单个 ABD 模型在其特定的“主场”配置下可以达到或超过固定块专家的水平。例如，在 LM1B 上，经过广泛分布训练的 ABD 在块大小为 4 和 16 时，达到了与专门为这些尺寸训练的 BD3LM 专家相当的困惑度。
零样本泛化：
- 在留出基准测试（PTB, Wikitext, Lambada 等）上，ABD 展示了鲁棒的泛化能力。具体而言，块大小为 4 的 ABD 在科学领域（Pubmed, Arxiv）的表现优于 BD3LM 专家，这表明多尺度训练提高了领域偏移的鲁棒性。
生成质量：
- 在生成困惑度方面（通过 GPT2-Large 测量），ABD 在同等计算成本 (NFEs) 下达到了最强固定块专家的质量水平，证实了广泛训练并不会降低样本质量。

消融研究

分布敏感性： $\pi$ 的选择允许从业者在不同尺度间进行权衡。指数权重倾向于小块；均匀权重倾向于大块；对数正态权重则提供平衡。
训练动态： 多尺度训练起到了稳定的正则化作用，能在提高所有块大小性能的同时，实现同步提升。
离散 vs 连续： 研究发现，相对于连续分布，基于 2 的幂次的离散分布更易处理且表现更好，这可能是由于它们代表了不同的操作点。

5. 重要性与主张

论文声称，配置支持集是决定扩散语言模型泛化能力的中心因素。通过将关注点从固定的架构约束转向训练配置的分布，ABD 建立了一个原则性的替代方案，用以取代固定的结构化目标。

解决不匹配： ABD 成功消除了训练-推理不匹配，允许单个模型处理任意解码策略（滑动窗口、自适应块），且无需架构改动或推理时的启发式方法。
统一模型： 它证明了单一模型可以实现结构不变性，在保持所有解码尺度下高质量的同时，恢复了块大小与困惑度之间预期的单调关系。
局限性： 作者谦虚地指出，性能仍然取决于 $\pi$ 的选择（它决定了尺度间的权衡），并且推理效率（在测试时选择最优策略）仍是一个开放性问题。其理论保证依赖于支持覆盖度，并未提供有限样本界限。

总之，自适应块扩散提供了一个稳健的框架，用于训练对解码配置具有不变性的扩散语言模型，解决了当前结构化扩散方法中的一个关键局限。