扩散语言模型：从架构挑战到推理优化的深度探索

★

原文链接：https://zhuanlan.zhihu.com/p/1998418717743289472

作者：王云鹤

在思考这一主题时，我首先回想起多年前一位前辈提出的问题：Transformer的下一代范式是什么？ 我当时认为，Transformer是长期技术积累从量变到质变的成果，其思想（如Non-local模块）在早期视觉研究中已有体现，而卷积与注意力机制也始终在互补发展。扩散模型本身并非Transformer的直接替代品，但其建模方式有潜力对自回归范式构成显著冲击。

我关注扩散语言模型已久，但受限于精力与算力，未能进行深度探索。从纯文本角度切入扩散模型架构是当前相对可行的路径，且其中诸多基础问题若不解决，多模态版本的进展也将受阻。因此，本文将首先聚焦于扩散语言模型的算法基础。

自去年下半年起，我们团队在若干方向上展开了探索。受一位内部专家启发，我在元旦前完成了一篇洞察性文章。近期在AAAI的报告中，我们重点介绍了团队的几项工作，包括Next-Block Diffusion训练、Diffusion in Diffusion分层架构以及Diffusion Agent等。

扩散语言模型：从架构挑战到推理优化的深度探索

相关演示文稿已发布在项目网站上（我始终认为，将扩散模型应用于代码与智能体场景极具前景）：
* 诺亚Diffusion LLM项目网站：https://noah-dllm.github.io/

关于当前扩散模型面临的挑战，我们也撰写了一份研究提案。其思想源于此前在伦敦、加拿大等地的一系列研讨，后经我整理为以下中文观点：

推理高效的架构（KV Cache / Attention / 模型架构）：当前扩散模型整体仍沿用自回归模型的框架，尤其是注意力机制（包括QKV计算与KV Cache）。在自回归的Next Token Prediction推理模式下，KV Cache的复用能极大提升效率。然而，扩散模型中掩码位置的随机性导致KV Cache复用机制失效，这是阻碍其广泛应用的核心瓶颈之一。因此，我们需要为扩散模型设计更适配的注意力结构或更具结构性的掩码方式，在保留其解码优势的同时，提升推理效率。
- 初步思路：从自回归的KV Cache机制出发，若能使扩散模型中的掩码Token在去噪过程中严格遵循从左到右的顺序，似乎可解决复用问题。但这并非完美方案，我们将尝试探索更专为扩散模型设计的注意力或KV相关机制。
更适配的词表（Tokenizer）：理想的扩散模型不应完全遵循自回归的现有范式，而应具备如人类思考般的结构性，例如先为给定主题生成全局大纲，再细化局部内容。本质上，人类在写作或报告时往往是多尺度的。但自回归模型的编码器通常基于统一的切分与统计算法，导致词元粒度单一。对于扩散模型，其Tokenizer是否也应具备结构性？例如，采用不同粒度的Tokenizer，有的负责段落间关联，有的负责细节修改，有的适合快速填充空缺。
- 初步思路：一个简单的设想是构建具有重叠词元的金字塔式词表，依据数据的全局信息（如首尾部分）进行统计。可参考图像生成中“先全局后局部”或反之的范式。但要发挥此类词表的优势，对应的训练与推理范式也需调整。
更好的优化范式（训练与优化）：扩散模型常被讨论的优化问题在于梯度计算效率较低。极端情况下，一个长序列（如128K）中仅有一个掩码Token，却需计算全部前向与反向传播，即为单个Token付出了大量计算，而获得的梯度反馈可能很小。这意味着同等训练开销下，扩散模型的精度往往不及自回归模型。此外，当前扩散模型在预训练阶段对数据施加随机掩码，而在指令微调阶段则对答案部分全部掩码，可能存在不同训练阶段不一致的问题，对后续的强化学习也构成挑战。掩码比例本身也是一个需要结合不同训练阶段进行动态调整的重要研究课题。
- 初步思路：这是一个极具挑战性的问题。目前能想到的是在训练过程中探索更高效的方案。同时，在预训练与指令微调阶段，可以引入更多融合策略以提升扩散模型的训练效果。
更优的掩码方式（多掩码Token）：当前主流扩散模型通常只使用一个掩码Token，这虽简洁有效，但多样性不足。所有被掩码位置均使用相同Token处理，在功能性上有所局限。此外，现有掩码范式中所有位置被处理的概率相等，一些简单位置可能无需掩码，且掩码位置间缺乏联动性，没有结构化的掩码机制。因此，如何改进当前较为初级的掩码方式，进一步挖掘扩散模型的潜力，是一个值得深入探讨的问题。
- 初步思路：将单一的掩码Token扩展为多个，并赋予其一定的先验假设，使新增的Token能够联动并高效工作。也可以考虑更先进的、尤其是结构化的掩码添加方式，这在代码等复杂任务场景中可能具有更大潜力。
动态长度输出：扩散模型虽具备良好的并行解码特性，能针对给定问题快速生成输出，但通常需要预先指定输出长度，并依赖EOS Token来终止生成。这类似于命题作文：要求写一篇800字的文章，扩散模型能在800字内终止，但难以处理超出长度的情况。尽管训练数据包含不同长度，使模型对预期长度有一定适应性，但在极端情况下（如比较9.11与9.8的大小却生成100K个Token）则显得低效且可能影响精度。因此，如何根据输入问题自适应推断最优输出长度，是一个值得探索的方向。
- 初步思路：在模型训练中同时加入对EOS位置的预测，使推理过程能感知输入问题，减少无效计算。此外，参数复用等技术可能对超出给定长度的外推任务有所帮助。
适配扩散模型的数据工程：当前大多数扩散模型复用了自回归模型的数据集。即便是由自回归模型改造而来的扩散模型，也继续使用其积累的数据与技术。这些通用数据蕴含的知识固然可被扩散模型学习，但若想进一步激发扩散模型的潜力，使其学习更多结构化知识并具备更强推理能力，现有数据仍有提升空间。由于扩散模型训练阶段的掩码Token是随机加入的，想要结合此特性在数据层面进行优化，面临的挑战和工作量依然巨大。
资源高效的模型优化：当前阶段，在扩散模型中探讨细粒度的模型压缩（如权重或神经元层面的稀疏化）可能为时尚早，因为许多基础模型架构尚未收敛。然而，提升扩散模型整体推理效率的研究价值显著，尤其是在批量处理（batch size）增大后，全局扩散推理在部分场景下相较于自回归模型呈现劣势。此外，提升去噪过程本身的高效性，对于任何形态的扩散模型都能带来显著收益。从现有评测看，自回归模型仍具明显优势，因此，如何协同使用自回归与扩散模型以获取更大的整体收益，也值得深入探讨。
- 优化思路：可优先尝试几个潜在收益较大的方向：扩散多步蒸馏（参考图像生成领域的相关工作，以减少去噪步骤）、投机推理、低比特后量化以降低推理开销。也可探索互补性的高效组合，例如在长序列、单批次场景中使用扩散模型，而在大批次信息整合场景中使用自回归模型。
慢思考及隐式思考：在扩散模型的监督微调过程中，对于给定的查询，模型会在预定的长度空间内通过迭代去噪生成答案并进行对比。若采用传统的顺序思维链方法来实现“慢思考”，对扩散模型而言可能效率不高，未能充分释放其潜力。在数据构造和优化方式上仍有很大改进空间。此外，扩散模型的另一大潜力在于支持“重掩码”（remasking）——在去噪过程中，可对置信度较低的标记进行重置和修改，这为深度思考和隐式思考提供了更多可能性。充分利用这些特性，有望进一步提升模型智能。
- 优化思路：首先，可能需要针对扩散模型的特性构建新的思维链形式，例如结构化的思维链。在代码生成和智能体场景中，从生成大纲框架入手，再补充细节，可能产生更优结果。其次，可对不同粒度的思考过程、过程中的编辑与重掩码操作进行更多探索。
结构化的提示词工程与记忆：尽管现有算法可通过将自回归模型续训至扩散模式，或通过块解码机制赋予扩散模型一定的局部自回归特性，但扩散模型的掩码与解码方式终究不同。其独特之处在于不仅能“向前看”标记，也能“向后看”。这引发了一个思考：是否存在更适配扩散模式的提示词格式与更高效的提示方法？例如，仅提供少数全局关键标记，即可快速驱动完整的解码与推理过程，这对代码生成、深度研究、智能体等场景极具价值。
- 优化思路：一种直接的方法是将传统的自回归问答式提示转变为适用于扩散模型的完形填空式提示，尤其可通过提供关键信息来辅助生成更佳回复。也可探索提示词自演进的方进行深度优化。同时，这些方法也可应用于检索增强生成和记忆相关的工作中。
未来的统一架构：在多模态或多视角任务中，一个端到端、从头训练的方案常被认为更能激发模型能力上限。然而，现实常面临诸多挑战：不同模态数据的对齐、数据配比，甚至不同任务所使用的模型与优化目标截然不同（例如，理解类任务当前多采用自回归模型，而生成类任务则多用扩散模型）。未来属于多模态，因此探索更统一的模型结构与训练范式至关重要，而扩散模型目前展现的能力正具备这种潜力。
- 优化思路：以视觉语言动作模型为例，这是一个热门领域，尤其在真实物理环境交互中已成功融合多种信息。当前主流架构中，视觉与语言部分多以自回归模式为主，而动作部分则采用传统（非离散掩码）的扩散模式以生成连续的动作轨迹。那么，能否使用离散扩散模型将这三个部分统一融合？这是一个非常有趣且值得探索的问题。