多块扩散语言模型 (Multi-Block Diffusion Language Models)

TL;DR MBD-LM 将块扩散语言模型从单块扩展至多块并行解码，通过多块教师强制训练-推理对齐，将平均每步生成 token 数从 3.47 提升至 6.19，而精度不降反升。

问题

问题背景

扩散语言模型（DLM）因生成质量与可控性潜力受到关注，但推理效率是其核心瓶颈。Block Diffusion Language Models (BD-LMs) 引入 KV 缓存与灵活长度生成，将文本划分为块进行单块扩散解码，但块间顺序执行限制了并行潜力。

现有方法局限

• 训练与推理不匹配：主流 BD-LMs 采用 Teacher Forcing 训练，每次仅暴露一个噪声块并以干净前缀为条件。这与多块推理（MultiBD）中同时解码多个连续块、各块噪声级别各异、前缀动态变化的运行集状态差异显著。
• Diffusion Forcing 的差距：虽引入多块噪声可见性，但其训练状态仍非 MultiBD 推理的有界运行集（bounded running-set），槽位噪声模式是异构的，且多块块间依赖关系未被建模。这导致推理时性能下降或需额外矫正。
• 解码实现僵硬：现有方案依赖动态输入形状，破坏 前缀缓存 复用，计算图需频繁重编译，并行度提升无法转化为实际加速（wall-clock acceleration）。

为什么这个问题难/重要

• 训练分布对齐：需让模型在训练中接触推理时的噪声分布——多块同时去噪、块级噪声调度差异大，且活跃块数量动态变化，对齐训练与推理状态是统计学习与计算图的双重挑战。
• 解码系统优化：要保持静态输入形状以复用缓存、避免计算图膨胀，同时支持块级并行推进与选择性更新，这需要精巧的 Block Buffer 机制设计。
• 业界关注度：提升扩散模型文本生成的 Tokens Per Forward pass (TPF) 直接关系到部署成本与时延，可类比 Speculative Decoding 追求更高的单步吞吐，但路径完全不同，对扩散模型生态有通用价值。

行业类比

这类似于从逐 token 的自回归解码升级为半自回归块解码，试图在扩散范式下实现类似 并行解码 的加速，但需克服训练与推理的鸿沟。

核心洞察

• **Multi-block Teacher Forcing** 通过 **bounded noise-groups** 和 **chain-uniform noise scheduling** 弥合了训练与推理的分布差异，使模型能有效处理 MultiBD 推理中的异构噪声状态。与单块 teacher forcing（仅见单块）或 diffusion forcing（全序列可见但训练噪声模式不匹配 bounded running-set）不同，MultiTF 直接模拟推理时的 slot-wise 噪声多样性，从而在提升并行度的同时保持甚至提高准确率。
• **Block Buffer 解码机制** 将多块并行解码从理论加速转化为实际 wall-clock 提升。它通过固定形状的块缓存、前缀 KV 缓存复用和静态输入形状，克服了动态 running-set 导致的形状重编译与缓存失效问题。实验表明，MBD-LLaDA2-Mini 的 TPF 从 3.47 增至 6.19，结合 DMax 可达 9.34，且 math/code 准确率仅下降 1.02%，为大模型扩散解码的工程部署提供了可行方案。

方法

输入：已训练的单块扩散语言模型

Multi-Block Diffusion Language Models（MBD-LMs）以一个预训练完成的 Block Diffusion Language Model (BD-LM) 作为起点，例如 LLaDA2-Mini。该模型原本仅在 单块扩散（SingleBD） 模式下运作，即每次只对一个干净前缀之后的单个噪声块进行去噪，解码并行度受限于块尺寸。

关键模块 1：多块教师强制（Multi-block Teacher Forcing, MultiTF）训练

为支持多块并行解码，MBD-LMs 采用 MultiTF 对原模型进行后训练，使模型适应“在一个滑动窗口中同时处理多个连续噪声块”的状态。训练过程构造了一个有界噪声组（bounded noise-group）：从序列中选取一段连续的 token 组，记为运行集（running-set），组内 token 被加上不同程度的噪声，而组前的 token 保持干净并作为条件前缀。具体包含三个设计：

• 组布局（group-layout）：随机选择组起始位置与长度，模拟推理时滑动窗口的不同偏移。
• 链式均匀噪声调度器（chain-uniform noise-scheduler）：对组内每个位置独立采样一个噪声时间步，但遵循“越靠后的块噪声越大”的单调趋势，以匹配推理时块间异构的噪声模式。
• 组感知双流掩码（Group-Aware Dual-Stream Mask）：在注意力计算中，干净前缀与噪声组、以及噪声组内不同时间步之间采用不同的可见性策略，确保训练与推理的因果逻辑一致。 训练目标沿用标准交叉熵，但对不同位置按噪声强度加权，引导模型学会从多块部分去噪的中间状态预测原始干净 token。

关键模块 2：块缓冲（Block Buffer）静态解码算法

推理时，朴素的多块扩散因序列长度动态变化而无法复用前缀缓存。MBD-LMs 提出 Block Buffer 机制：

• 维护一个固定大小的块窗口（由 K 个块组成），窗口内 token 构成静态形状的输入张量。
• 每轮去噪后，根据置信度阈值更新各块 token，并将已解码完成的块移出窗口，同时在末尾装入新的噪声块，窗口整体前移。
• 由于缓冲区形状固定，前缀 KV 缓存可被逐轮累积复用，大幅减少自回归部分的重计算。
• 解码循环控制由阈值门控实现：仅当某块所有 token 置信度超过阈值才视为完成，否则继续去噪。

输出与效果

MBD-LMs 将单块解码扩展为多块并发，每个前向传递可并行生成多个块的 token，提升吞吐量。例如，MBD-LLaDA2-Mini 将平均每步 token 数从 3.47 提升至 6.19，同时准确率从 79.95% 微升至 81.03%。结合 DMax 增强解码后，TPF 可达 9.34，而数学与代码基准准确率仅下降 1.02%。

与同类方法的差异

相较于先前扩散语言模型中试图引入多块可见性的 Diffusion Forcing (D2F)，MultiTF 的关键区别在于：它直接对齐有界运行集的推理分布，通过随机化组布局与链式均匀噪声调度，覆盖推理时所有可能的块偏移和噪声组合，而非仅训练无界的、不切实际的完整序列扩散轨迹。这使得后训练出的模型无需复杂 adaptor 即可直接部署于高效的多块并行解码。

实验

实验基于 BD-LM 架构（LLaDA2-Mini），通过 MultiTF 后训练从单块扩散扩展到多块扩散。评估指标包括每前向通行 token 数（TPF）和下游任务准确率，覆盖数学和代码推理基准。设计对比了 SingleBD、MultiBD 及集成 DMax 的去噪策略。

关键发现：MultiTF 训练有效对齐了多块推理的动态噪声模式，MBD-LLaDA2-Mini 的 TPF 从 3.47 提升到 6.19，准确率反而从 79.95% 升至 81.03%。这在提升并行度的同时未牺牲质量。结合 DMax 的激进去噪策略，TPF 进一步达到 9.34，准确率仅小幅下降 1.02%，显示了推理加速与质量间的可调节权衡。

深度解读：多块扩散的核心挑战在于训练-推理分布偏移——传统 teacher forcing 仅见单噪声块，而 MultiTF 通过有界噪声组和链式均匀调度弥合了这一鸿沟。与单纯的扩散强制相比，该方法更贴近实际推理的滑动窗口状态，且通过 Block Buffer 保持静态输入形状和前缀缓存复用，使得理论并行度切实转化为墙钟加速。这为扩散语言模型的实用部署提供了新范式。

行业影响

落地场景

多块扩散语言模型（MBD-LM） 通过并行解码多个文本块，大幅提升生成速度，适合对延迟敏感的实时生成场景：

• 对话与客服系统：聊天机器人、智能客服需要快速响应，MBD-LM 可将一次前向传递生成的 token 数从 3.47 提升至 6.19，结合 DMax 可达 9.34，显著降低用户感知延迟。
• 代码补全与编程助手：代码生成通常具有结构化块特性，多块并行生成可以同时产出多行代码，提升交互流畅度；实验在数学与代码基准上仅损失 1.02% 准确率，实用性高。
• 内容创作平台：新闻摘要、产品描述等批量生成服务，MBD-LM 在保持质量的同时，提高吞吐量，降低单位生成成本。

商业价值

• 推理成本大幅降低：更高的 Tokens Per Forward pass (TPF) 意味着相同文本量所需的 model forward 次数减少，直接降低 GPU 时租成本。对于大规模在线服务，可节省 30%–50% 的推理算力。
• 用户体验提升：端到端延迟缩短，尤其在交互式应用中，响应速度直接影响留存率和转化率。
• 兼容现有自回归流程：MBD-LM 的后训练方案无需从头预训练，可在已有 BD-LM 基座上微调，迁移成本低，加速产品落地。

与现有产品/工作流的接口

• 即插即用的解码器替换：MBD-LM 提供的 Block Buffer 机制 保持静态输入形状和前缀缓存复用，可直接集成到基于 Transformer 的推理框架（如 vLLM、HuggingFace TGI）中，无需重构服务架构。
• 训练-推理对齐：多块教师强制（MultiTF） 在微调阶段便模拟多块推理的噪声模式，避免了训练-推理不一致问题，模型导入后即可获得加速效果。
• 兼容增强解码策略：可与 DMax 等 token 级优化技术无缝叠加，进一步提效，无需单独调整。

具体用例

1. 电商商品描述批量生成：大促期间需生成数百万条个性化商品文案。采用 MBD-LM，每条描述生成延迟减半，相同 GPU 集群下吞吐翻倍，直接降低云服务成本，同时保证文案质量。
2. 在线教育 AI 助教：实时解答学生疑问，要求 200ms 内开始输出。MBD-LM 的并行块解码可在首 token 延迟不变的情况下加速后续 token 生成，避免学生等待，提升学习体验。

局限

• **训练仍然依赖已有的 BD-LM 基础模型。** MBD-LMs 本质上是对预训练好的 Block Diffusion Language Model 进行后训练（post-training）得到，无法从零训起。尽管 MultiTF 缓解了训练-推理差异，但推理时采用动态 running-set 与训练时固定大小噪声组仍然存在偏差，尤其当解码过程中 active set 尺寸变化或块边界不匹配时，可能导致解码质量波动。此外，MultiTF 的训练目标（如 chain-uniform scheduling）增加了调参负担，对不同下游任务可能需重新适配。
• **泛化性受限于所测试的模型架构和解码策略。** 论文实验主要围绕 LLaDA2 系列模型开展，虽然在 SDAR 等少数 BD-LM 上进行了验证，但尚未覆盖连续空间扩散模型或非自回归架构的变体。同时，MultiBD 解码的性能增益与 block size、running-set 大小等超参强相关，面对不同任务（如极其短文本或非结构化生成）时可能无法保持同样显著的速度-质量平衡。
• **Block Buffer 设计牺牲了一定的灵活性。** 为保持静态输入形状和缓存复用，Block Buffer 要求块大小、最大 running-set 窗口在推理前固定，当生成长度变化剧烈或需要动态调整并行粒度时，该方法可能次优。此外，现有实现中块间仍保留因果约束（block-causal），虽较逐 token 解码有提升，但并未完全释放非自回归的并行潜力；当模型规模增大或序列极长时，前缀缓存的内存开销也可能成为瓶颈。