大语言扩散模型来了，何必只预测下一个token | 人大高瓴&蚂蚁

明敏 发自 凹非寺
量子位 | 公众号 QbitAI

用扩散模型替代自回归，大模型的逆诅咒有解了!

人大高瓴人工智能研究院、蚂蚁共同提出LLaDA（a Large Language Diffusion with mAsking）。

LLaDA-8B在上下文学习方面与LLaMA3-8B能力相当，而且在反转诗歌任务中超越GPT-4o。

在大语言模型领域，反转诗歌是一个特殊任务，它用来评估模型在处理语言模型的双向依赖关系和逻辑推理能力。

比如让大模型写出“一行白鹭上青天”的上一句。

通常情况，自回归模型（如GPT）根据下文推断上文的表现上总是不够好。这是因为自回归模型的原理就是利用序列中前面的元素来预测当前元素，即预测下一个token。

而LLaDA是基于扩散模型的双向模型，天然能够更好捕捉文本的双向依赖关系。

作者在摘要中表示：LLaDA挑战了LLMs关键能力与自回归模型之间的固有联系。

这些研究也引发了不少讨论。

有人提出：

我们正在重构掩码语言模型建模？

RAG和嵌入式相似性搜索上，这种范式也可能表现更好？

值得一提的是，LLaDA仅用13万H800GPU时训练了2.3万亿token语料，然后对450万对token进行SFT。

正向掩码+反向预测

论文核心提出了一个问题：自回归是否是实现LLMs智能的唯一路径？

毕竟自回归范式的LLMs目前仍存在诸多弊端，比如逐个生成token的机制导致计算成果很高，从左到右建模限制了逆推理任务中的性能。

这都限制了了LLMs处理更长、更复杂任务的能力。

为此，他们提出了LLaDA。通过正向掩码和反向预测机制，让模型更好捕捉文本的双向依赖关系。

研究采用标准的数据准备、预训练、监督微调（SFT）和评估流程，将LLaDA扩展到80亿参数。

在2.3万亿token上从零开始预训练，使用13万H800 GPU时，随后在450万对数据上进行监督微调。

在语言理解、数学、代码和中文等多样化任务中，表现如下：

强大可扩展性：LLaDA 能够有效扩展到10²³ FLOPs计算资源上，在六个任务（例如MMLU和GSM8K）上，与在相同数据上训练的自建自回归基线模型结果相当。

上下文学习：值得注意的是，LLaDA-8B 在几乎所有 15 个标准的零样本/少样本学习任务上都超越了 LLaMA2-7B，并且与 LLaMA3-8B表现相当。

指令遵循：LLaDA在SFT后显著增强了指令遵循能力，这在多轮对话等案例研究中得到了展示。

反转推理：LLaDA有效地打破了反转诅咒，在正向和反转任务上表现一致。特别是在反转诗歌完成任务中，LLaDA 的表现优于 GPT-4o。

LLaDA使用Transformer架构作为掩码预测器。与自回归模型不同，LLaDA的transformer不使用因果掩码（Causal Mask），因此它可以同时看到输入序列中的所有token。

模型参数量与传统大语言模型（如GPT）相当，但架构细节（如多头注意力的设置）略有不同，以适应掩码预测任务。

正向掩码过程如下：

LLaDA采用随机掩码机制，对一个输入序列x0，模型会随机选择一定比例的标记进行掩码（masking），生成部分掩码的序列xt。

每个token被掩码的概率为t，其中t是从[0,1]中均匀采样的。这与传统的固定掩码比例（如BERT中的15%）不同，LLaDA的随机掩码机制在大规模数据上表现出更好的性能。

模型的目标是学习一个掩码预测器，能够根据部分掩码的序列xt预测出被掩码的token。训练时，模型只对被掩码的token计算损失。

其中1[·]是指示函数，表示只对被掩码的token计算损失。

在SFT阶段，LLaDA使用监督数据（如对话对、指令-响应对）进一步优化模型，使其在特定任务上表现更好。

对于每个任务，模型会根据任务数据的特点进行微调。例如在对话生成式任务中，模型会学习如何根据给定对话历史生成合适响应。

在SFT阶段，模型会根据任务数据的特点选择性地掩码响应部分token，这使得模型能够更好地学习任务相关的模式。

推理部分，在生成任务中，LLaDA通过反向采样过程生成文本。从一个完全掩码的序列开始，逐步预测出被掩码的token，直到生成完整的文本

采样过程中，LLaDA采用多种策略（如随机重掩码、低置信度重掩码、半自回归重掩码）来平衡生成效率和质量。

在条件概率评估任务中，LLaDA会根据给定的提示（prompt）和部分掩码的响应（response）来评估模型的条件概率。这使得LLaDA能够在各种基准任务上进行性能评估。

预训练LLM在不同基准上的表现如下。

后训练后在不同benchmark上的表现如下。其中LLaDA只进行了SFT，其他模型有进行额外的强化学习对齐。

在反转诗歌任务中，LLaDA超越了GPT-4o。

在多轮对话任务中LLaDA的表现如下，较深颜色表示采样后期阶段预测的token，较浅颜色表示在采样早期预测的token。

网友：期待能被真正用起来

研究团队同时放出了一些LLaDA的实际表现。

可以解决普通的数学推理问题。

编程问题也OK。

有国外网友表示：这肯定会推动中国AI研究更加关注小模型。不过也不代表他们放弃scaling。

同时也有人评价说，这或许可以开启一些混合模型的可能。

以及有人提及Meta也有过类似的工作，将transformer和diffusion相结合。

当然也有人关心，此前也提出了不少超越Transformer的架构，但是它们都还没有被学术界/工业界真正采纳。

让我们期待后续吧。

本项研究由人大高瓴人工智能学院与蚂蚁集团共同带来。通讯作者为李崇轩，他目前为人大高瓴人工智能学院准聘副教授，目前focuses的方向为深度生成模型，了解现有模型的能力和局限，设计有效且可扩展的下一代模型。

论文地址：
https://arxiv.org/abs/2502.09992
项目主页：
https://ml-gsai.github.io/LLaDA-demo/