这是一篇尝试改变LLM「范式」的文章:当前主流的LLM架构都是「自回归」的,通俗地理解就是必须「从左到右依次生成」。这篇文章挑战了这一范式,探索扩散模型在 LLMs 上的可行性,通过随机掩码 - 预测的逆向思维,让模型学会「全局思考」。
论文:[2502.09992] Large Language Diffusion Models
背景
主流大语言模型架构:自回归模型 (Autoregressive LLMs)
过去几年,自回归模型(Autoregressive Models, ARMs)一直是大语言模型(LLM)的主流架构。典型的自回归语言模型以Transformer解码器为基础,按照从左到右的顺序依次预测下一个词元(token)。
形式化地,自回归模型将一个长度为\(N\)的文本序列\(X=(x_1, x_2, ..., x_N)\)的概率分解为各位置的条件概率连乘积:
这种“下一词预测”范式在GPT-2、GPT-3、PaLM、LLaMA等模型中取得了巨大成功。它经过大规模数据的最大似然训练后,能够生成流畅连贯的文本,并表现出惊人的Few-shot学习和上下文推理能力
然而,一个根本性的问题一直存在:“难道只有自回归范式才能实现LLM的智能吗?” 。近年来的实践表明,自回归模型虽然有效,但并非完美无缺。在推进LLM性能的过程中,人们逐渐发现了自回归结构的一些局限性:
- 计算效率问题: 自回归模型在推理时无法并行生成。由于每个新词元的生成依赖前面的所有词元,必须严格按顺序逐字生成,这导致推理速度随输出长度线性降低。对于长文本或实时应用,逐字生成的延迟成为瓶颈。例如,生成一段长度为\(n\)的文本需要\(n\)次前向推理,无法像图像生成那样一次性产生所有像素。这种序列生成开销高。
- 长程依赖与逆向推理难题: 自回归模型的单向性意味着它只在一个方向(通常为左到右)建模文本依赖关系。当任务需要反向推理或利用后文信息时,AR模型往往力不从心。一种典型情况是所谓“逆向诅咒 (Reversal Curse)”——如果模型在训练中只见过"A 是 B"类型的语句,那么让它生成"B 是 A"就会遇到困难。研究发现,即使对于简单的可逆关系,传统LLM对调输入顺序后性能会急剧下降。这是因为自回归模型训练时始终假定左侧为条件、右侧为目标,缺乏从右向左推理的能力。
- 上下文利用的限制: AR模型虽然可以利用数千词的上下文,但只能将未来信息纳入隐式表示,无法直接条件在后续词上。这对一些任务不利,例如诗歌或代码补全:模型在生成时如果前面缺少关键信息,可能无法正确“回看”未来需要填的内容。而双向模型(如BERT)在训练时可同时看左右上下文,但传统BERT并非生成模型,无法自然生成文本。
- 其他局限: 有研究将AR语言模型等价为无损压缩器;尽管这显示出LLM压缩知识的能力,但本质上任何足够复杂的生成模型(不局限于AR)也能实现类似功能。同时,AR模型在非常长文本的全局一致性、重复和幻觉等问题上也表现出一定劣势,需要更复杂的注意力机制或后处理来缓解。这些都激发了研究者思考:是否存在不同于自回归的方法,能够避免这些局限?
扩散模型的成功及启示
在图像、音频等连续信号领域,近年兴起的扩散模型 (Diffusion Models)为生成建模提供了新的可能。扩散模型通过逐步添加噪声再去噪声的过程来生成数据,在图像生成任务上取得了革命性突破
扩散模型成功的关键在于:它将生成过程划分为多步,每一步进行小幅修改,而非一次性生成全部内容。具体而言,扩散模型包含正向过程(逐步向数据添加噪声,使其趋于随机分布)和反向过程(逐步去除噪声,从纯噪声恢复数据分布)。在图像领域,这种逐步生成方式带来了稳定训练和高样本质量的优点。
扩散模型的潜力对NLP领域具有启发意义:文本虽然是离散符号序列,但理论上也可以设计类似的逐步“噪声扰动-还原”生成过程。如果这种思路应用得当,可能在以下方面优于自回归模型:
- 并行生成: 扩散的反向生成可在每个步骤中同时更新多个位置的词元,从而有希望并行化部分生成过程,减少生成时延。
- 双向依赖: 反向去噪过程可以利用全局上下文(部分词元已填充,部分待生成),模型在预测时能够同时参考已生成内容的左侧和右侧。这天然缓解了自回归模型不能利用未来词的缺陷,使模型对顺序倒置的任务更为健壮。
- 新的优化目标: 扩散模型训练基于对数似然的下界(ELBO)优化,这与直接的交叉熵训练不同,可能提供不同的正则效果。理论上,任何符合生成建模基本原理(最大似然)的模型都应具备LLM的通用能力。因此,用不同于AR的方式来逼近语言分布,也有望学习到相似的能力。
LLaDA
基础概念
在介绍核心方法之前,我们先梳理生成建模和两类模型范式的基本概念,包括生成概率建模原理、自回归 vs. 扩散的区别,以及掩码扩散模型的原理。这将有助于理解LLaDA的设计哲学和数学推导。
生成建模基本原理
最大似然估计 (Maximum Likelihood Estimation, MLE): 所有大语言模型本质上都是生成式模型,旨在对训练语料的真实但未知分布\(p_{\text{data}}(x)\)进行近似建模。给定带参数\(\theta\)的模型分布\(p_{\theta}(x)\),我们希望\(p_{\theta}\)尽可能接近\(p_{\text{data}}\)。为此,通常采用最大似然训练,相当于最小化数据分布与模型分布之间的KL散度:
\( \theta^* = \arg\max_{\theta} \mathbb{E}_{x \sim p_{\text{data}}} [\log p_{\theta}(x)] \)
等价于
最大似然使得模型在训练数据上概率最大,也即模型分布逐步逼近真实分布。对于离散文本数据,MLE训练通常实现为交叉熵损失,梯度为负对数似然。
自回归建模: 如前所述,自回归模型通过链式法则分解联合概率。在MLE框架下,自回归模型的训练目标是:
这正是标准的因果语言模型训练的交叉熵损失(将每个词元预测看作一个分类问题累加)。AR模型的优势在于这一参数化天然满足归一化且易于采样(一次生成一个,下一个条件在前面已生成的)。它直接优化对数似然,因此最大似然原则下,该模型能够在极限情况下逼近真实分布。许多LLM的能力(如few-shot、指令遵循)被认为来源于这种因果预测训练。
扩散建模: 扩散模型属于隐式生成模型,它不直接给出\(P_{\theta}(x)\)的解析表达,而是定义一个分步生成过程,通过变分下界来优化似然。以经典连续扩散为例,有一系列连续状态\({x_t: t\in[0,1]}\),\(x_1\)是纯噪声,\(x_0\)是原始数据。正向过程\(q(x_{t}|x_{t-1})\)逐步加入噪声,反向模型\(p_{\theta}(x_{t-1}|x_t)\)学习去噪。模型训练目标是ELBO(Evidence Lower Bound),形式为各步的KL散度之和:
再加上\(q(x_T|x_0)\)与先验\(p(x_T)\)的KL。通过优化此下界,模型逐步逼近真实数据的对数似然。虽然扩散模型的训练目标不像AR那样直观(它将原始log-likelihood分配到多个子目标上),但理论保证这一目标也是对\(\log p_{\theta}(x_0)\)的下界,优化它等价于提高模型似然。因此,只要模型有足够表达能力并在海量数据上训练,扩散模型原则上也能学到和AR模型相媲美的生成能力。
KL散度方向: 需要注意,自回归MLE对应最小化\(\text{KL}(p_{\text{data}} | p_{\theta})\),倾向于覆盖数据分布(避免遗漏模式,但可能生成某些低概率样本);而扩散模型ELBO分解后,有些变体等价于某种\(\text{KL}(p_{\theta} | q)\)或混合形式。这可能导致行为上的细微差异,例如扩散模型往往生成更多样化的样本。这在DiffuSeq等工作中体现为较高的生成多样性。
自回归 vs. 扩散:核心区别
- 依赖建模方式: AR模型是单步的、单向的。给定前文,它一次生成下一个token,完全由前文决定,无法直接调整已经生成的内容。相反,扩散模型是多步的、迭代的。它通过多轮推理逐渐完善序列,每一步都可以利用目前序列中的所有已知部分。这意味着扩散模型能建模双向依赖:因为在生成过程的中间状态,一个token的生成可以参考其左侧和右侧已有内容。这种双向性类似BERT等遮盖语言模型在训练时看到双向上下文,但扩散模型是真正生成式的,在推理时同样利用双向信息。
- 并行化潜力: 在理论上,自回归生成需要\(N\)步才能得到\(N\)个词元,而扩散模型若设计得当,可以在\(T \ll N\)的步骤内同时确定多个词元的内容。例如,如果每一步“揭示”若干个掩码,\(T\)步后完成序列,则可以比\(N\)步更快。极端情况下,若能在一步内填满所有掩码,就实现了完全并行的生成(当然往往需要多步权衡质量)。因此扩散模型提供了效率与质量的折中:步数多则质量高、但慢;步数少则快、但可能损失精度。
- 训练难度与稳定性: AR模型训练相对简单稳定(交叉熵损失,容易收敛),而扩散模型训练需要平衡各扩散步骤的目标,涉及更复杂的技巧。例如,在离散扩散中,如果步骤太多或噪声增加策略不当,模型可能难以训练或在采样时遇到数值问题(如概率校准不良)。近期有工作指出很多掩码扩散模型实际上对时间步不敏感,本质更接近无序的遮盖模型,需要注意采样时的准确性 。为此,有研究引入连续时间公式、重参数化、重要性采样等手段改进训练稳定性。简而言之,扩散模型的实现复杂度更高,需要仔细设计掩码策略和目标函数以取得良好效果。
- 模型输出灵活性: AR模型每步输出一个token的概率分布,最终只能得到一个确定序列。而扩散模型的每一步输出是对整个序列的更新,本质上可以纠正先前步骤的错误。在实际应用中,这意味着扩散模型更适合交互式生成或逐步细化场景。例如,在某些对话或写作辅助场景下,模型可以初步生成草稿,再通过扩散过程 refine 文本以满足特定要求(这类似人类写作的过程)。AR模型若要实现类似效果,需借助后处理或编辑模型。扩散模型天然地内置了多次调整机会,这赋予它灵活性。
- 理论表达能力: 从理论上说,只要生成过程能覆盖数据空间,AR和扩散都具备充分表达能力。但有趣的是,Ou等人(2024)证明,离散扩散模型的损失上界可以解释为某种任意顺序自回归模型的期望损失。也就是说,一个扩散模型的训练等价于在各种可能顺序下预测下一个token的损失平均。这提示两者在理论上并非截然不同,而是相关联的:扩散模型相当于综合了多种生成顺序的信息。这也侧面解释了为何一个良好训练的扩散模型应当能具备自回归LLM的诸多能力——因为它并未违背生成模型的基本原则 (最大似然等),只是实现路径不同。
Masked Diffusion Models, MDM
针对离散文本序列,当前应用最成功的扩散策略是掩蔽扩散 (Masked Diffusion),又称吸收扩散。其核心思想是:将“噪声”定义为将词替换为掩码符号[MASK]。正向过程逐渐把原文本的词元随机替换成[MASK],直到全序列被掩盖;反向过程则从全掩码开始,不断预测恢复出原来的词。具体原理如下:
- 吸收态马尔可夫链: 设初始数据序列为\(x_0 = (w_1, w_2, ..., w_N)\),这里\(w_i\)表示第\(i\)个token(词或子词)。定义特殊掩码符号
[MASK],代表未知/被噪声覆盖的状态。正向扩散\(q(x_t | x_{t-1})\)在每一步以一定概率将未掩盖的token变为[MASK],而已经是[MASK]的token保持为[MASK](吸收性)。经过若干步\(T\)后,\(x_T\)将趋近于全[MASK]序列。例如,在简单设置下,可令每个token在每步有固定概率\(\beta\)被掩盖,则\(T\)步后每个token被掩盖的总体概率接近1。这个过程类似不断向序列添加“噪声”(即掩码),最终破坏所有信息。 - 反向预测过程: 反向过程\(p_{\theta}(x_{t-1} | x_t)\)由掩码预测模型 (Mask Predictor)实现。掩码预测器以当前序列\(x_t\)为输入,尝试填补其中所有[MASK]位置的真实内容。由于前向过程保证\(x_t\)中未被掩盖的token在\(t-1\)时刻就是正确的原始token(未被掩盖过),反向过程在这些位置不需要改动,只需预测[MASK]处应是哪一个token。这等价于让模型在给定部分单词的上下文下同时预测多个空缺。在完全理想的情况下,经过一次预测就可以恢复到\(t-1\)的状态;但实际上模型是概率预测,每次只能近似采样出可能的填充,需要多步迭代逼近真实分布。
- 训练目标: 掩码扩散模型通常使用交叉熵损失训练,使得掩码预测器在给定任意部分掩盖的序列时,都能正确复原出原序列。公式上,训练采样一个随机时间步\(t\)(对应一定比例的掩码)以及对应的损坏序列\(x_t\),然后对\(x_t\)中的每个[MASK]预测其原始值\(x_0\),计算预测概率的对数损失:\[ \mathcal{L}_{\text{mask}}(\theta) = \mathbb{E}_{x_0 \sim \text{data}} \; \mathbb{E}_{t \sim \mathcal{U}(1,T)} \; \mathbb{E}_{x_t \sim q(x_t|x_0)} \Big[ -\sum_{i: x_{t,i}=\texttt{[MASK]}} \log P_{\theta}(x_{0,i} \mid x_t) \Big] \tag{5}\]
上式中,仅对那些在\(x_t\)中被掩盖的位置计算损失,其余位置的token已揭示,不计入损失。这种训练类似BERT的掩码语言模型,但有两点关键区别:
- 掩码比例\(t\)是动态变化的,可以从0到100%随机取值;
- 最终目标并非学习语义表示用于下游,而是直接用于生成,因此要求模型对任意比例的掩码都能很好重建。
- 对数似然下界解释: 尽管表面上只是交叉熵预测掩码,实际上可以证明,上述损失函数是模型最终产生数据分布的负对数似然的一个上界。也就是说,最小化\(\mathcal{L}_{\text{mask}}\)等价于最大化模型生成序列的似然下界,使其成为一个原理上有效的生成模型训练目标。Ou等人2024年的理论分析进一步表明,这个损失事实上等价于扩散模型ELBO并产生Fisher一致性(无偏估计)。因此,掩码扩散模型并非一种启发式训练,而是有坚实的概率基础支撑的。
- 掩码比例调度: MDM的一个关键要素是如何选择各步的掩码概率/比例。在训练中,可以简化为直接随机掩盖一定比例的tokens,而不真正模拟一步步掩盖的过程。例如,训练时随机采样掩盖率\(m\sim [0,1]\),然后将输入序列按此比例随机掩盖,喂给模型预测。这相当于从整个扩散链中采样一个中间状态\(x_t\)进行训练。这种方法被称为“一跳训练”,类似连续扩散中随机采样一个噪声水平训练去噪网络,能够高效地覆盖所有扩散阶段。同样,掩盖概率可以采用线性调度(从0逐渐到1),也可以根据公式计算。但实践中Uniform随机足以覆盖各种情况。重要的是,掩码比例取0对应无噪声,取1对应全掩盖,这两个极端确保模型既能学习恒等映射(0%掩码,不需改动)也能学习从无到有的生成(100%掩码,需生成完整序列)。 Masked Diffusion与BERT的区别: 表面上,掩码扩散的训练和BERT掩码语言模型相似,都是随机掩盖部分token并预测 。但有以下显著不同:
- BERT使用固定的掩码率(如15%),而MDM使用随机可变的掩码率(0-100%均匀)。这意味着MDM训练覆盖了从极少掩码到几乎全掩码的全谱情况,迫使模型学习在各种已知/未知比例下都能推断缺失内容。
- BERT的目标在于学习深层双向表示用于下游任务,本身并非生成模型,即使训练完成也不能自由采样新句子(因为它没有定义序列整体概率)。MDM则明确以生成为目的,可通过多次迭代采样生成新句子。可以说,MDM是让BERT学会如何自己把一句话“拼出来”,而不仅是理解填空。
- BERT仅进行一次掩码预测就结束(单步静态填空),而MDM将在生成时多次重复掩码->预测过程。尽管训练阶段是一跳的,但推理时需要把这些能力串起来,逐步从100%掩码走到0%掩码。为此,生成时要设定扩散步数\(T\)并采取一定的调度(如均匀间隔)在这些步上应用模型。因而,MDM训练出的模型可以看作“一次掩码预测器”,需要封装在一个多步算法中才能作为生成模型使用。 掌握以上概念后,我们可以深入探讨LLaDA提出的大语言扩散模型如何具体运作。接下来首先介绍LLaDA的方法框架和算法推导,然后阐述其实验结果与能力分析。
LLaDA核心思路和算法
LLaDA整体框架概览
LLaDA遵循当前大模型的标准训练流程:首先大规模无监督预训练,然后进行有监督微调(SFT)来注入指令遵循等能力。区别在于,预训练阶段LLaDA并不是优化自回归的预测损失,而是优化掩码扩散模型的损失。因此,其模型结构也略有不同:LLaDA的核心是一个Transformer掩码预测器,而非通常的因果Transformer解码器。不过在大部分配置上(参数规模、数据集、优化器等),LLaDA和同规模的自回归LLM保持一致,以保证可比性
整个框架可以总结如下:
- 数据准备: 收集和清洗大规模语料,构建语料库(包含中英文、代码、数学公式等多领域文本)。LLaDA使用了约2.3万亿token的数据,和LLaMA等模型的数据量级相当。分词器选用与GPT系列类似的BPE分词(词表略有不同,因为数据分布不同)。
- 预训练阶段: 使用掩码扩散策略对Transformer进行训练。每个训练样本(序列)随机选取一个掩码比例,将该序列中对应比例的token替换为[MASK],然后让模型预测这些被掩码的位置原本的token。损失只计算预测错误的掩码token部分。通过随机选择不同的掩码比例,模型学习在从几乎无缺失到几乎全缺失的各种场景下重建文本。这种训练持续若干轮遍历直至收敛。
- 监督微调阶段: 将模型适配为能够遵循指令的人机对话模型。具体方法是使用成对的(prompt, response)数据进行有监督训练。对每个对话样本,保持prompt部分不变,将response部分按一定掩码策略进行掩盖,然后让模型在有prompt上下文的条件下预测response掩码处的内容。LLaDA的方法是仅掩盖回答的token(prompt保留),这相当于模型学习条件生成:在给定提示下生成回答。损失同样是对掩码部分计算交叉熵。经过此阶段,模型具备了遵循自然语言指令的能力,例如问答、多轮对话等。
- 推理阶段: 给定用户输入(可以视作prompt),LLaDA以扩散采样的方式生成输出:首先确定输出的长度(或最大长度),初始化一个全掩码序列,然后迭代地应用掩码预测器来逐步产生完整回答。在每一步,模型根据当前已填充的部分预测所有[MASK]标记应替换成的词,采样或选择概率最高的词填入,得到一个更少掩码的新序列。然后可以重新随机选取一部分词再次掩码(称为re-masking策略,用于保持多样性和避免一次决策失误),再交由模型预测。如此循环进行\(T\)步,直到没有掩码或达到设定步数。输出即为模型的生成文本。这一过程与图像扩散中的逐步去噪类似,只不过这里“噪声”是未填充的[MASK]标记而已。 在LLaDA框架下,预训练模型(Base)和指令微调模型(Instruct)都是通过掩码扩散实现生成的。因此,它们在推理时都不使用常规的自回归生成方法,而是使用上述掩码迭代填充算法。这一点对直觉来说比较新颖:以往LLM无论预训练还是SFT后,推理基本都是左到右逐词输出,而LLaDA的推理更像在玩“逐字填词游戏”。但正是这种生成方式,让LLaDA克服了一些AR模型的局限。下面我们深入解析其数学模型和算法细节。
LLaDA数学模型推导
前向掩码过程定义: 考虑长度为\(N\)的离散序列\(x_0=(x_0^1,\dots,x_0^N)\)表示原始文本(这里\(x_0^i\)是第\(i\)个token的值)。定义一个吸收马尔可夫链\({x_t}*{t=0}^T\),其中\(x_T\)趋向于全掩码序列。形式化地,我们可以设定一个序列\({m_t}*{t=1}^T\)表示掩码概率调度,即在前向过程中,从\(t-1\)到\(t\)步,每个token会被掩码的概率。典型地,可以取\(m_t\)递增,例如\(m_t = \frac{t}{T}\)(线性增加)或者其他设计。
则前向扩散\(q(x_t|x_{t-1})\)可定义为对每个token独立地: \( q(x_t^i = \texttt{[MASK]} \mid x_{t-1}^i = w) = m_t, \qquad q(x_t^i = w \mid x_{t-1}^i = w) = 1 - m_t,\) \(q(x_t^i = \texttt{[MASK]} \mid x_{t-1}^i = \texttt{[MASK]}) = 1,\) 即:如果前一时刻token是具体词\(w\),则它有概率\(m_t\)变为掩码,有概率\((1-m_t)\)保持原值;如果前一时刻已经是掩码,则保持掩码(吸收性)。经过\(t=1\)到\(T\)步后,\(x_T\)中的每个位置非掩码的概率约为\(\prod_{s=1}^T (1-m_s)\),若\(m_s\)设计得足够大,\(x_T\)几乎肯定全为掩码(即\([MASK],[MASK],...\))。这完成了将原始数据逐渐“噪声化”的过程。
值得注意的是,由于吸收性,前向过程的逆条件分布\(q(x_{t-1}|x_t, x_0)\)非常简单:在\(x_t\)已知和原始\(x_0\)已知的情况下,\(x_{t-1}\)除了那些在\(t\)时刻仍然未掩盖的token必然相同,其余位置就等于\(x_0\)对应token。这种结构被Austin等人称为结构化离散扩散 。
反向生成过程定义: LLaDA引入参数\(\theta\)的掩码预测模型\(p_{\theta}(x_{t-1} | x_t)\)来近似\(q(x_{t-1}|x_t, x_0)\),也就是学习在不知道\(x_0\)的情况下,从\(x_t\)重构\(x_{t-1}\)。由于\(x_t\)相比\(x_{t-1}\)只是多了一些掩码,我们可以把\(p_{\theta}\)进一步拆解为对每个掩码位置进行预测。具体而言: \( p_{\theta}(x_{t-1} = u \mid x_t) = \prod_{i=1}^N p_{\theta}(x_{t-1}^i = u^i \mid x_t),\) 其中\(u\)是一个可能的\(N\)长度序列,在\(x_t\)未掩码的位置\(u^i\)必须等于\(x_t^i\)(因为这些位置已经有值,不应改变),在\(x_t\)掩码的位置\(u^i\)可以为词表中的任何词。这一分布由Transformer掩码预测器计算。换句话说,模型读取当前序列\(x_t\)(包含若干[MASK]标记和若干已知词),并输出每个[MASK]应替换成的词的概率分布。由于Transformer注意力机制可以全局查看序列,这些预测利用了序列的双向上下文。
优化目标(ELBO): 训练的目标是让\(p_{\theta}(x_{t-1}|x_t)\)尽可能接近真实的后验\(q(x_{t-1}|x_t, x_0)\)对于所有\(t\)和所有\(x_0\)。这可以通过最大化下界(ELBO)实现:
LLaDA的训练正是基于此下界,对每一项KL进行采样近似并最小化。幸运的是,由于\(q(x_{t-1}|x_t,x_0)\)仅在\(x_t\)掩码的位置与\(x_0\)不同(否则两者相同),上式展开正得到前述的掩码预测交叉熵损失形式 (5)。也就是说,最小化掩码预测损失(5)即在最大化模型对数似然下界。因此,LLaDA训练遵循严格的概率模型准则,确保训练过程收敛后模型的生成质量有保证。
训练算法实现: 实际训练中,LLaDA并不显式模拟每一步掩码,而是采取一跳采样:直接采样一个\(t\)并构造\(x_t\)。算法流程如下:
- 对于每个训练批次样本:取原始序列\(x_0\);
- 随机采样一个时间步\(t \sim \text{Uniform}({1,\dots,T})\)(若采用离散步数),或者直接采样一个掩码率\(\gamma \in [0,1]\)(连续情形,将其视为\(t\));
- 根据\(t\)或\(\gamma\),以对应概率将\(x_0\)的每个token替换为
[MASK],得到损坏序列\(\tilde{x}\),作为\(x_t\)的一个样本; - 将\(\tilde{x}\)输入模型,得到对每个掩码位置的预测分布,与原始\(x_0\)进行比较,计算损失;
- 对该损失关于\(\theta\)求梯度,更新模型参数(使用AdamW或其它优化器)。
通过上述过程,模型被迫学习在随机mask数量的情况下复原序列。尤其地,当\(t\)非常接近\(T\)时,序列几乎全掩盖,模型需要学习从几乎无上下文的情况下凭空生成(这类似语言模型的常规生成);当\(t\)很小(接近0掩盖)时,模型主要学习保持已给定词不变。这两端都很重要:前者赋予模型生成能力,后者确保模型不会改变不需要改变的上下文(这有助于保持一致性和记忆提示)。这一训练策略与Nie等人2024年的建议一致,即加入变长序列训练来增强模型泛化。在实现中,LLaDA甚至让1%的训练样本随机选取一个短于正常长度的子序列训练(动态长度),以确保模型能适应不固定长度的输入输出情境,这对于处理对话等场景很有用。
模型参数与结构: LLaDA的掩码预测器采用标准Transformer架构,与同规模LLM非常相似。例如,LLaDA-8B模型使用了与LLaMA3-8B相当的层数和隐藏维度。区别在于:LLaDA的Transformer不使用因果mask。也就是说,多头注意力机制可以让每个token的表示看见整个序列,包括其右侧的token。这样设计是因为在掩码预测任务中,没有生成顺序的概念,所有位置并行预测,因此不需要遮盖未来信息。换句话说,LLaDA的骨干更接近BERT这样的双向Transformer,但参数规模和许多细节又对应GPT类模型。例如,LLaDA取消了LLaMA3中使用的分组查询注意力(GQA)机制,改回标准多头注意力,以简化实现。这样做的原因在于扩散模型不需要Key-Value缓存(因为不是自回归逐步预测),如果用GQA在没有缓存的情况下反而浪费参数,所以直接用密集注意力并调整前馈层宽度以平衡参数量。
总体来说,LLaDA模型结构与主流解码器Transformer区别不大,只是注意力不掩盖未来和没有缓存优化。这意味着训练和推理时计算量略有增加(每一步都需全序列计算注意力),但也带来了表示能力更强的好处(可利用全局信息)。
具体示例:句子补全任务中的工作方式

为了形象说明LLaDA的生成过程,我们以一个简单的句子补全例子来演示其工作流程。假设我们有一句话:“机器学习是人工智能的一个分支”,希望模型生成其下一句。对于自回归模型(如GPT),典型做法是将这句话作为提示直接续写;而对于LLaDA,由于模型不是按顺序生成,而是基于掩码扩散,我们需要设置问题为一个条件生成任务:
- 构造输入: 将用户提示(上一句)作为prompt,接着准备一个空的响应区域。例如输入序列可以表示为:“机器学习是人工智能的一个分支。
[MASK]…”,其中句号后面是一串连续的[MASK]标记,长度为我们预估的回答长度(比如128个[MASK]表示生成128字的续写)。这是LLaDA推理的起点\(X_T\),即全掩码的待生成部分。Prompt部分在SFT后的模型中被视为条件,不会被掩盖。 - 设置步数和调度: 决定扩散的总步数\(T\)以及每步掩码替换多少。在推理中,我们通常选定一个固定的步数,例如\(T=100\)步,并采用线性解码,即每一步解除一定比例的掩码,使最后一步恰好无掩码。简单来说,可以每步揭示\(1%\)的掩码,这样100步后全部填完。当然实际实现中更灵活,LLaDA作者默认用均匀时间步调度,每步试图去掉固定数量的掩码。他们指出步数越多,结果越好,但越慢;步数减少可以加快生成,但略损失质量,存在折中可调。
- 迭代生成:
- 第1步 (\(t=T\) 到\( t=T-1\)): 当前序列是全掩码(只有prompt非掩码)。模型看到prompt和一长串
[MASK],会尝试为所有这些\(M\)(mask)位置预测可能的词。由于完全没有上下文,模型的预测主要依据prompt提供的信息和训练中学到的一般语言分布。在实践中,可能会引入一个策略:并非直接填写所有掩码,而是一次只填一部分,或随机填一部分、留下一部分掩码保持未定。这种策略称为局部采样或重掩盖(remasking)。例如,LLaDA可能每步只填充其中的2个token,然后将它们视为确定,其余仍保持掩码用于下步预测。这样做可以避免一次性决策过多造成错误累积。第1步后,序列也许变成:“机器学习是人工智能的一个分支。它[MASK]...[MASK]”,其中“It/它”这个词被填入了第一个掩码位置。 - 中间若干步: 模型反复读取新的部分填充序列,然后预测剩余掩码。例如,在第2步,模型已知续写的开头是“它”,结合prompt前文,可能推测整句想表达的是定义或性质,因此继续填充下几个词,如“是当前
[MASK]领域[MASK]”等等。每一步模型都利用先前生成的词作为上下文,有点类似人类逐句写作先想好大致结构再逐词完善。随着步数增加,掩码逐渐减少,序列信息越来越完整,模型预测也越来越受具体上下文约束。 - 最后步 (t=0): 理想情况下,经过\(T\)步后,所有掩码都被具体的词填充,得到最终生成的完整句子。例如最后模型生成结果为:“机器学习是人工智能的一个分支。它是当前计算机科学领域中发展最快的方向之一。”(假设模型这么续写)。此时\(t=0\)对应完全去噪文本。这个输出就是LLaDA给出的回答。
- 第1步 (\(t=T\) 到\( t=T-1\)): 当前序列是全掩码(只有prompt非掩码)。模型看到prompt和一长串
- 结果输出: 模型输出的序列(去除了掩码)即为对用户提示的续写回答。与自回归模型不同的是,这个过程并非严格从左到右,而是在全句范围内逐步细化。如果在某些步产生了不满意或不一致的词,后续步还有机会通过重新掩码并预测来纠正。比如模型可能在第10步时写了“不正确”的词,但后续觉得上下文不协调可以将其再次掩盖替换。事实上,LLaDA支持一种半自回归 (semi-AR) 重掩盖策略:每步总是将刚填的最后几个词掩盖一部分重新预测,以确保衔接流畅。这些技巧都利用了扩散模型能反复修改的特点。 通过以上示例可以看到,LLaDA的生成类似多人合作填字游戏:模型每次“填写”一点,然后检查整体是否通顺,再调整未定部分。这种迭代生成方式带来的直接好处是:模型可以处理任意顺序的依赖。例如,如果后面的内容需要前面信息才能确定,模型可以先暂时跳过,待前文生成后再回来补全。这有效克服了AR模型“只能顺着来,倒着不行”的问题 。正如LLaDA作者所示,在需要从后文推导前文的逆向生成任务上,它表现出色。下面我们详细讨论这一点。
克服“逆向诅咒”和增强上下文学习
逆向诅咒 (Reversal Curse): 先前提到,自回归LLM如果训练数据大多是顺序“一览”,在遇到需要逆向推理的场景时会遇到困难。一个典型实验是:给模型一对相关句子(A, B),让它在看过“A是B”后生成“B是A”。Berglund等人(2023)指出GPT-3等LLM在这种任务上失败率高,称之为“reversal curse”。因为它们在训练时几乎从未尝试倒过来生成句子。
LLaDA通过掩码扩散天然缓解了这一问题:在LLaDA的生成过程中,序列的任意部分都有可能在后面的步骤才生成,这相当于模型学会了“晚到的信息也能补充前面的空白”。例如,在掩码扩散里,句子前半部分和后半部分谁先生成是不确定的;有时模型可能先填充了句尾的一些词,之后再来完善句首的细节。这种非因果的生成顺序练就了模型的逆向预测能力。
论文中一个有力的实验证明了这点:他们构建了496对中国诗句(一句出自诗的上联,另一句是下联),让模型在顺向(给上联预测下联)和逆向(给下联预测上联)两种情况下进行零样本生成。结果显示,LLaDA在两种方向上的生成性能一致,几乎同样擅长根据上文生成下文和根据下文生成上文。换言之,它成功打破了“只能一向”的魔咒。而对比的自回归模型(如Qwen-2.5和GPT-4系列)在逆向任务上表现明显较差,即顺向生成可能完美,但逆向就错漏百出。这种巨大差距正是传统LLM逆向诅咒的体现。LLaDA则有效消除了顺/逆任务间的性能落差。在某些案例中,LLaDA甚至超过了GPT-4的表现,例如在诗句补全任务中取得更高的合理性评分。
上下文学习能力: 上下文学习(In-Context Learning, ICL)指LLM在给定一些示例的上下文中,能类比地完成新任务,而无需参数更新。自回归LLM具备惊人的ICL能力,一度被认为与其自回归结构相关(因为它把上下文示例串联,然后预测下一个,与训练形式一致)。然而,LLaDA的研究表明,这种能力其实是通用生成模型的一种内在属性。LLaDA虽然不是自回归形式,但在零样本/少样本学习任务上,同样表现出色,甚至在一些基准上超越了同规模的自回归模型。例如,LLaDA-8B在15个标准Few-shot任务上几乎全面超过自回归的LLaMA2-7B,并与最新的LLaMA3-8B持平。这说明,只要模型学到了条件分布(无论通过哪种机制),就能利用提示中的示例来调整生成分布,产生正确输出。因此,ICL并非AR模型的独占优势,而是任何符合生成原则的LLM都可以具备的。
更强的条件推理: 扩散模型的双向性还带来另一个优势:模型可以更灵活地利用提示。在指令微调后,LLaDA在多轮对话等场景展现出优秀的指令遵循和上下文理解能力。例如,人们观察到LLaDA产生的回复在引用上文内容、遵循格式要求等方面都很到位。这可能归功于掩码预测器在预训练时已经习惯了“看全局再发言”。因此,在微调指令数据时,它更容易学会遵从指令和正确引用提示信息。
此外,由于生成可以跳跃进行,LLaDA在需要长程规划的生成任务上也具备潜力。例如,写代码、证明题等任务往往需要先有大致思路再具体展开。LLaDA可以先产出提纲性的内容,然后在后续步骤细化;而AR模型要么提前规划再逐字输出(这很难直接学会),要么就一边生成一边想,很可能走偏再难以纠正。理论上,如果配合一些启发式策略,扩散式LLM在复杂推理任务上可能更有优势。
模型优化目标和训练策略细节
接下来,我们补充LLaDA在训练实施中的一些重要细节,这些策略保障了模型的性能和稳定训练:
- 变长度训练: LLaDA在1%的预训练数据上采用了随机截断,即随机选择一个小于最大长度的值\(L\),截断序列长度为\(L\)进行掩码训练。这样做是为了让模型学会处理不同长度的文本,尤其是在序列结尾生成终止符的情况。这对实际应用很关键,因为模型应对短句和长文都能自适应。这一点在自回归LLM中也常用,但对扩散模型尤为重要——因为全掩码生成时需要知道生成何时停止或输出何种长度。
- 重掩盖 (Remasking) 策略: 为了进一步提高生成质量和稳定性,作者在推理采样中采用了一种半自回归的重掩盖策略。具体做法是:在每一步填充之后,并不直接前进到下一步,而是可能将新填的token的一部分再掩盖掉,留待后面的步骤重新预测。这样避免了一次决策的失误永久保留下去。例如,他们的一个方案是设置采样步数为生成长度的一半,确保平均每步引入2个新token。这意味着,每生成的两个词中,就有一个词可能在后续被替换调整。这样既保证最终长度正确,又给了模型纠错的余地。实验表明,这种机制能提升困难任务(如数学题)的准确度。
- 采样步数的影响: 作者对采样步数\(T\)对性能的影响做了分析。在GSM8K数学问答和HumanEval代码生成上测试发现,当步数从较低增加到与输出长度相当时,性能明显提升。这符合预期:步数越多,逐步生成越细致,错误可以慢慢纠正。但超过一定步数后收益递减,而且计算成本直线上升。因此在实际使用中,需要根据任务复杂度和响应长度选择合适的步数。例如,他们对指令微调模型通常设置步数与回答长度相等,从而保证最佳质量;而在一些消融实验中,为了效率使用了步数=长度/2以加快测试。总之,步数是一个质量-效率调节杆,可根据需求调节生成速度。
- 无监督指导 (CFG): 类比图像扩散中的Classifier-Free Guidance,作者也探索了为LLaDA引入指导信号的方法。他们在附录提到,可以对Prompt条件分别生成有无提示的两种预测,计算差异以提升针对提示的一致性(这类似于图像扩散里放大条件的影响)。不过由于计算资源限制,正式评测里并未使用这种指导技巧。这是一个未来可提升的方向:也许通过一定的推理技巧,LLaDA可以进一步提高对提示的服从度或生成所需风格(如安全/风趣)。但需要注意过强的指导可能影响生成流畅度,需要权衡。
- 优化器和学习率策略: LLaDA预训练使用AdamW优化器 (权重衰减0.1),总batch约1280,分布式训练。学习率采用Warmup-Stable-Decay方案:先热身升至峰值(迭代2000步),然后长时间保持恒定,训练过约1.2万亿token后开始线性衰减,再在最后0.3万亿token线性降低到很小。这一策略确保前期稳定训练,中期充分学习,后期缓慢收敛,有助于大模型训练的收敛稳定。LLaDA-8B总计训练了2.3T token,花费约13万 H800 GPU小时。训练过程中并未报告明显不稳定或发散问题,可见这样的超大规模扩散模型训练在工程上是可行的。 总的来说,LLaDA通过周详的模型设计和训练策略,成功在千亿参数、万亿token规模上训练了掩码扩散模型,使其在能力上足以比肩甚至超越传统自回归LLM。接下来,我们将深入介绍实验评测结果和结论,包括LLaDA在各类基准上的表现、与现有模型的比较以及模型的优势与局限。
实验和结论
LLaDA论文通过一系列严格的实验评估,证明了大语言扩散模型的可行性和竞争力。本节概述主要实验设置和结果,包括训练细节、基准测试表现、与对比模型的比较,以及对LLaDA可扩展性和局限的讨论。
训练配置与资源
模型规模: 作者训练了两个主要配置的模型:LLaDA-1B和LLaDA-8B,参数规模分别约1亿和80亿。8B是主力模型,用于和其它8B级别LLM比较。模型层数、隐藏维度等超参数与同等参数量的自回归Transformer相匹配,以保证可比性。
预训练数据: LLaDA预训练使用2.3万亿token的海量语料。数据来源包括通用网络抓取文本(类似Common Crawl,但经过严格清洗)、高质量百科和书籍、代码语料库(开源代码)、数学公式与学术论文,以及中英文多语种数据。为保证多样性并覆盖不同任务,数据混合比例参考了已发布的大模型(如LLaMA2等)经验,但也结合了一些自构建AR模型的小规模试验确定混合策略。低质量内容通过规则和LLM过滤剔除,以提高训练效率。序列长度统一设为4096,每个训练样本截断/填充到该长度。与LLaMA等相比,这个长度偏长,保证了LLaDA能处理长上下文。
计算资源: 训练在NVIDIA H800 GPU集群上完成,总计约0.13百万(13万)GPU小时。这个规模与训练同尺寸AR模型所需计算相当。说明在大致相同算力投入下,可以训练出扩散式LLM达到类似效果。8B模型训练仅进行了一次完整运行,未进行多次超参调优。这意味着报告结果并非精挑细选,是模型自然性能的体现。
微调数据: 监督微调(SFT)使用了约450万条指令响应对。这些数据涵盖多种域和任务,包括:
- 自然语言指令及回答(如开放问答、解释型回答)、
- 代码生成与调试对话、
- 数学题解答过程、
- 表格和结构化数据问答、
- 多轮对话 (对话式QA,chat) 等。
数据来源可能包括公开的English和中文的指令数据集、开源对话数据、代码问答库,以及作者自建的一些数据(如PoE平台或ShareGPT对话转写)。他们遵循与已有LLM相同的微调协议(如OpenAI的InstructGPT方法或LLaMA等的SFT流程),未使用强化学习对齐(RLHF)等更高级的对齐技术。这一点很重要:LLaDA只是用最基本的SFT就达到了很强的指令跟随性能arxiv.org,没有额外借助人类反馈优化,显示了其固有能力。
评测方法: 为了公平比较,作者在评测时构建了一系列自回归基线模型。这些基线与LLaDA共享相同的数据和训练步骤,只是将预训练改为标准的因果语言建模(即GPT式训练)。尤其是,他们训练了一个同为8B参数的自回归Transformer模型,称为“ARM-8B”(可视为与LLaDA-8B直接可比的GPT式模型),在2.3T token上训练并在450万指令上SFT。另一个对比是引用了Meta最新的LLaMA3-8B模型。LLaMA3据报道是Meta的下一代模型,在架构上有一些改进(如Grouped Query Attention)。作者尽量消除了这些超参上的差异,使得比较主要反映训练范式(AR vs Diffusion)的影响。
基准任务表现
论文在多方面评测了LLaDA模型。

与现有模型的对比分析
- 性能对比: 综合以上结果,LLaDA-8B的性能全面可比于同规模的自回归LLM,如LLaMA2-7B、LLaMA3-8B等。在一些任务上甚至略有超越,如few-shot分类和翻译任务上LLaDA表现尤为突出。这说明掩码扩散模型在同等资源投入下并不逊色,打破了业界对“自回归才有效”的偏见。相比更大的模型(如Claude 2 细化版、PaLM2等,几十B以上),LLaDA-8B还有明显差距。但作者认为这主要是因为参数规模和训练数据的限制——8B毕竟属于小模型,在一些复杂常识推理和编码任务上无法与数百亿参数的模型竞争。这种差距并非扩散方法导致,而是尚未把扩散模型推向更大规模进行验证。
- 可扩展性: LLaDA团队特别关注扩散模型的可扩展性。他们指出,过去许多人怀疑扩散模型难以扩展到海量数据和参数,但LLaDA的结果证明了扩展性良好。例如,从1B提升到8B参数,LLaDA在各项任务上大幅提高,很多任务的提升幅度与自回归模型的“scaling law”一致。这表明Fisher一致性等理论属性在LLaDA中体现:随着模型更大、数据更多,其逼近真实分布的能力加强。因此没有出现因为方法不同而遇到的扩展瓶颈。作者甚至推测,如果给LLaDA上百亿甚至上千亿参数和更多训练语料,它完全有潜力达到GPT-4等最强模型的水准。当然,实现这一点需要巨大的算力投入和工程优化,但没有理论障碍。
- 训练成本: 从成本看,LLaDA-8B训练消耗与训练一个自回归8B模型相当。这意味着掩码扩散并没有带来额外的指数级计算增长。虽然一次推理采样需要多步,但训练中通过随机采样t已经降低了多步目标的冗余。同时,训练时可以充分并行(所有mask位置同时预测,相当于每个batch提供了比自回归更多的预测信号)。因此,训练效率不输AR模型。
- 推理成本: 推理方面,如果要求相同质量,自回归模型需要生成\(N\)个token做\(N\)次前向计算;扩散模型如果采用\(T\)步、每步并行决策若干token,则大致需要\(T\)次前向。\(T\)的典型值可能和\(N\)在同一量级(如用于长度512的回答可能\(T=256\)或512),所以推理开销目前略高于自回归模型。但是可以权衡:如果能容忍质量稍降,可用\(T \ll N\)来加速。另外,由于没有因果依赖,扩散模型也许更适合硬件优化并行,例如可以设计自定义芯片一次推多个步骤。此外,LLaDA不使用KV缓存,每一步都要重新计算整段注意力,这部分开销也可以通过改进注意力机制或缓存部分未变的信息来优化(近期NVIDIA等研究已经探索扩散Transformer的缓存加速)。所以推理效率虽目前稍弱,但有优化空间,不构成不可逾越的劣势。
独特优势: 除了上述性能本身,LLaDA还有一些AR模型无法轻易实现的优势:
- 天然的双向填空能力,使它可以做任意位置的插入/续写。例如,要在一篇文章中间插入一句话,自回归模型需要重写后续所有内容或使用复杂的插入LM,而LLaDA可以直接对中间空白位置掩码然后生成即可(本质就是其正常任务)。
- 一致的训练/推理对象: LLaDA在训练时就见过各种mask比例的情况,包括接近全mask(需要完全生成)和少mask(需要润色给定文本)。这意味着它在改写/纠错任务上也很有潜力——直接把需要改的字词mask掉,让模型重写即可。这比起要微调AR模型去做编辑要方便。
- 抗顺序偏差: AR模型有时会对序列的位置产生偏好(如序列越长错误积累越多)。扩散模型因多次校正,输出质量对长度变化更平稳。作者在附录分析不同长度的生成准确率,发现LLaDA长度从256增加到1024时准确率下降幅度小于自回归模型。
- 无需额外压缩的对齐: 由于LLaDA仅用SFT已经达到接近RLHF模型的效果,表明许多指标仅稍差,这意味着也许不需要复杂的RLHF,仅靠扩散模型内在的训练准则就能学到类似对齐的行为。这可能与扩散模型倾向\(\text{KL}(p_{\theta}|p_{\text{data}})\)有关,使其生成更靠近训练数据分布(而AR模型容易出现奇特输出,需要RLHF修正)。 局限性: 然而,LLaDA目前也存在一些局限和挑战:
- 模型规模较小: 8B参数在今天还不算顶尖模型,因而LLaDA未能在绝对性能上超越最强的封闭源模型(GPT-4、PaLM2等)。这并非方法的问题,而是实现上的阶段性。作者也承认需进一步扩大规模来全面评估扩散LLM的上限。
- 推理速度较慢: 如上讨论,扩散生成多步迭代,目前每一步要全序列计算,因而8B模型在长文本生成时速度比7B GPT慢。不过对于较短输出(几十字以内),扩散模型可以用较少步就完成,反而可能更快。这一点具体要视应用场景而定。
- 实现复杂度高: 相比AR LM一个简单的generate()接口即可采样,扩散LM需要实现掩码迭代逻辑,包括选步数、掩码策略、采样方法等。这增加了使用的复杂度和不确定性。对于非专家用户,上手可能没AR模型容易。目前社区也缺少成熟的扩散LM推理库,需要一定时间完善。
- 理论和直觉差异: 扩散LM的行为有时不如AR LM直观。例如,难以预测它何时会纠正前面错误,或输出长度如何收敛。因此调试和解释相比AR LM更具挑战。需要新的理论工具来分析其多步动态,已有一些研究在探讨这方面。
实验结论
综合实验结果,LLaDA取得了以下主要结论:
- 可扩展性验证: 扩散模型同样可以扩展到数十亿参数和万亿级语料,并展现出与AR模型类似的性能随规模提升趋势。这打破了人们对扩散模型“不好扩展”的疑虑
- 竞争性能: LLaDA-8B在各项基准上达到了与自回归LLM相当的水准,证明扩散模型是可行的LLM替代方案。在few-shot、逆向生成等方面甚至有独特优势
- 泛化能力: 扩散模型预训练的LLM同样具备上文学习、指令遵循、编程等能力,这些能力并非AR模型独有。生成建模原理本身赋予了模型这些通用智能特性。
- 逆向诅咒破解: LLaDA证明“顺序诅咒”不是不可避免,生成范式的改变即可解除这种限制。这为处理许多需要双向推理的任务开辟了新路。
- 未来工作方向: 作者指出需要进一步探索更大模型、多模态、采样加速等方向。扩散模型有潜力与视觉、音频模型结合,实现统一的多模态生成;以及通过算法改进提高推理速度,更实用化。
LLaDA2.0
训练过程--从 AR 到扩散模型的进化

LLaDA2.0 的训练并非从零开始,而是基于一个强大的 AR 基座模型(AR base model)。整个流程被设计为“知识继承”与“渐进式适应”的过程,通过三个递进阶段,让模型平滑地从传统的逐词生成模式,过渡到高效的非自回归(Non-Autoregressive)扩散生成模式。
- 第一阶段:持续预训练 (Continual Pre-training):从 AR 过渡到 MDLM(掩码扩散语言模型)
这是连接两种生成模式的桥梁。在此阶段,模型基于一个强大的AR Base model开始学习双向的、去噪式的生成方式。- 模型不再是从左到右预测下一个词,而是学习重建被随机掩盖(masked)的 token。
- 在保留原有模型表征几何结构(representational geometry)的同时,填补了自回归生成与扩散生成之间的鸿沟。
- 第二阶段:块扩散预训练 (Block Diffusion Pre-training):从 Token 级过渡到 Block 级建模
这是提升效率的关键步骤。在 MDLM 的基础上,训练进一步升级。- 模型不再处理单个 token,而是被训练去对连续的文本片段(即“Blocks”)进行去噪。
- 这种转变显著提高了计算效率,并增强了生成过程中的长距离连贯性(long-range coherence)。
- 第三阶段:后训练 (Post-training):对齐与任务专业化
当模型掌握了 Token 级和 Block 级的非自回归生成能力后,进入最后的打磨阶段。- 包含监督微调(SFT)和直接偏好优化(DPO)。
- 将扩散预训练中获得的强大生成能力转化为实际应用能力,确保模型遵循人类指令(Instruction-following),并符合下游任务的需求。
第一阶段:Continual Pre-training via WSD
这一阶段的核心策略是 Warmup-Stable-Decay (WSD)。它的目标是解决 AR 模型(从左到右)与扩散模型(双向、随机掩码)在架构偏置上的冲突,实现平滑迁移。
- Warmup (预热):渐进式块大小增加:将块大小 \(L_B\)从 1(AR 模式)逐步增加到 4096(全序列 MDLM 模式)。使用通用的 BDLM 损失函数。\[\mathcal{L}{\text{BDLM}}(\theta) = -\mathbb{E}{t, x_0, x_t} \left[ \frac{\alpha't}{1-\alpha_t} \sum{k=1}^K \sum_{i=1}^{L_B} \mathbb{1}[x_{t,k}^i = \text{[MASK]}] \log p_\theta(x_{0,k}^i | x_{0,<k}, x_{t,k}) \right] \tag{7}\]
公式表示模型在给定前文干净块 \(x_{0,<k}\) 和当前噪声块 \(x_{t,k}\) 的情况下,预测当前块中被掩码(MASK)的 Token。使得模型逐渐扩大模型的“感受野”,让模型从习惯于只看左边,慢慢适应双向上下文和联合去噪。
具体可以参考:离散扩散模型用于文本生成 - Stable (稳定):大规模 MDLM 训练:固定块大小 \(L_B = 4096\)(即\( K=1\),全序列作为一个块)。此时模型完全运行在 MDLM(Masked Diffusion LM)模式下。使用BDLM 损失的简化版(MDLM 损失):\[\mathcal{L}{\text{MDLM}}(\theta) = -\mathbb{E}{t, x_0, x_t} \left[ \frac{\alpha't}{1-\alpha_t} \sum{i=1}^{L} \mathbb{1}[x_{t}^i = \text{[MASK]}] \log p_\theta(x_{0}^i | x_{t}) \right] \tag{8}\]不再切分块,模型基于整个序列的噪声版本 \(x_t\) 来恢复被掩码的 Token。并利用大规模数据深入学习扩散动力学。由于不需要维护自回归的 KV 缓存逻辑,此阶段训练效率极高。
3. Decay (衰减):块大小回落
- 操作:将块大小 \(L_B\) 从 4096 逐步回调到较小的值(如 32)。
- 公式:回归使用 BDLM 损失函数(8)(同 Warmup 阶段)。
- 目的:知识蒸馏。将 MDLM 学到的全局上下文能力“压缩”回小块结构中。
- 意义:小块结构(BDLM)在推理时可以复用 KV-cache,从而实现快速推理。这一步是为了在“生成质量”和“推理速度”之间找到最佳平衡点。
第二阶段:监督微调 (SFT)
预训练让模型学会了“怎么生成文本”,SFT 则让模型学会“根据指令生成文本”。
核心改变:引入提示词(Prompt, \(c\))作为条件。使用条件化的 BDLM 损失。
其中,
- \(\mathbf{c}\):用户输入的指令/Prompt。
- \( x_{0,<k}\):之前已经生成的干净文本块(Auto-regressive context)。
- \(x_{t,k}\):当前正在去噪的噪声块。
训练模型在给定 Prompt \(c\) 的约束下,利用扩散过程生成符合预期的回复 \(x_0\)。赋予模型指令跟随(Instruction Following)的能力。
第三阶段:直接偏好优化 (DPO)
这一阶段是为了让模型的回复更符合人类的价值观(如有用性、安全性)。
核心挑战:扩散模型无法像 AR 模型那样轻易计算精确的对数似然(Log-likelihood),因此标准 DPO 公式不可用。
解决方案:使用 ELBO(证据下界) 近似替代对数似然。
- 定义条件 ELBO (\(B_{\text{BDLM}}\))
首先定义一个用于估算生成概率的项,本质上是 SFT 损失的负值(通过单次蒙特卡洛采样估计):
- DPO 优化目标
比较模型在“偏好回复”(\(x_w\)) 和“非偏好回复”(\(x_l\)) 上的 ELBO 差值
其中 \(\Delta B(x|c) = B_{\text{BDLM}}(\theta, x|c) - B_{\text{BDLM}}(\theta_{\text{ref}}, x|c)\),表示当前策略模型相对于参考模型的优势。
目的: 最大化好回复 (\(x_w\)) 与坏回复 (\(x_l\)) 之间的 ELBO 差距。
意义:在不需要训练独立奖励模型(Reward Model)的情况下,直接优化模型以符合人类偏好,解决了扩散模型在 RLHF 中的数学难题。