【LLaDA】Large Language Diffusion Models

note这篇论文介绍了LLaDA一种从头开始训练的大规模扩散语言模型。LLaDA展示了强大的可扩展性、上下文学习能力和指令跟随能力达到了与领先的LLMs相当的性能。此外LLaDA提供了双向建模和增强的鲁棒性有效解决了现有LLMs的一些局限性。LLaDA不是像 GPT/LLaMA 那样 从左到右一个 token 一个 token 生成而是先把句子里的 token 随机 mask 掉模型去预测这些 mask 的 token推理时从“全 mask”开始逐步把句子补出来LLM 的 in-context learning、instruction following、scaling这些能力是不是天然只能靠自回归模型得到作者的答案是不一定。LLaDA预训练时对整句按随机比例 maskSFT 时只 mask response生成时从全 mask 的 response 一步步还原。LLaDA 在 general、math、code、Chinese 等 benchmark 上整体能做到和他们自己训练的 ARM baseline 差不多8B base 还和 LLaMA3 8B Base 接近。SFT 后也能做多轮对话和指令跟随。AR 模型很擅长顺着生成但对“反着来”的任务天然不舒服。论文拿中文诗句补全举例给下句推上句LLaDA 因为是双向 mask 预测没有左到右偏置所以在 reversal task 上明显更稳甚至在这个任务上超过 GPT-4o。文章目录note一、研究背景二、LLaDA1. 前向数据掩码过程2. 反向生成过程3. 模型分布定义三、实验设计四、结果分析Reference一、研究背景研究问题这篇文章要解决的问题是大语言模型LLMs是否只能依赖于自回归模型ARMs来实现其核心能力如可扩展性、上下文学习和指令跟随。研究难点该问题的研究难点包括如何在不使用自回归模型的情况下通过生成模型原则实现LLMs的核心能力如何在有限的计算预算下实现大规模语言模型的扩展。相关工作该问题的研究相关工作有自回归模型在LLMs中的广泛应用和成功扩散模型在视觉数据上的成功应用以及对现有扩散模型在语言建模中的潜在扩展的研究。二、LLaDA这篇论文提出了LLaDALarge Language Diffusion with mAsking一种从预训练和监督微调SFT范式下从头开始训练的扩散模型。预训练时对整句按随机比例 maskSFT 时只 mask response生成时从全 mask 的 response 一步步还原。1. 前向数据掩码过程首先LLaDA通过逐步独立地掩码序列中的令牌来生成一个部分掩码的序列。对于时间t ∈ ( 0 , 1 ) t \in (0, 1)t∈(0,1)序列x t x_txt​是部分掩码的每个令牌被掩码的概率是t tt未被掩码的概率是1 − t 1 - t1−t。2. 反向生成过程然后LLaDA通过迭代预测掩码令牌来恢复数据分布。核心是一个掩码预测器一个参数模型p θ ( ⋅ ∣ x t ) p_\theta(\cdot|x_t)pθ​(⋅∣xt​)它接受x t x_txt​作为输入并预测所有掩码令牌。该模型使用交叉熵损失进行训练L ( θ ) ≜ − E t , x 0 , x t [ 1 t ∑ i 1 L 1 [ x t i M ] log ⁡ p θ ( x 0 i ∣ x t ) ] \mathcal{L}(\theta) \triangleq - \mathbb{E}_{t, x_0, x_t} \left[ \frac{1}{t} \sum_{i1}^L \mathbb{1}[x_t^i M] \log p_\theta(x_0^i|x_t) \right]L(θ)≜−Et,x0​,xt​​[t1​i1∑L​1[xti​M]logpθ​(x0i​∣xt​)]其中x 0 x_0x0​是训练样本t tt是从[ 0 , 1 ] [0, 1][0,1]均匀抽取的连续随机变量x t x_txt​是从前向过程中采样的L LL是序列长度3. 模型分布定义一旦训练完成可以通过掩码预测器参数化的反向过程来模拟反向过程并在t 0 t 0t0时定义模型分布p θ ( x 0 ) p_\theta(x_0)pθ​(x0​)作为边缘分布。损失函数在方程(4)中被证明是模型分布负对数似然的上界使其成为生成建模的有力目标。三、实验设计数据收集预训练语料库由来自公共来源的多样化数据构成包括网络数据、书籍、学术文章、社交媒体、百科全书、数学和代码约11%是中文61%是英文28%是代码。SFT数据集包括100万个人工标注样本和350万个合成样本。数据预处理数据清理涉及PDF文本提取、去重和有害内容过滤。为了确保质量使用BERT模型进行自动化数据质量注释以选择更高质量的样本。模型训练LLaDA采用Transformer架构作为掩码预测器并使用AdamW优化器和Warmup-Stable-Decay学习率调度器进行训练。预训练过程中使用固定序列长度4096计算成本为0.13百万H800 GPU小时。SFT过程中使用动态序列长度策略确保所有样本具有相同的长度。四、结果分析可扩展性LLaDA在六个标准任务上展示了令人印象深刻的可扩展性整体趋势与ARMs高度竞争。特别是在MMLU和GSM8K任务上LLaDA表现出更强的可扩展性。基准测试结果在预训练2.3T令牌后LLaDA 8B Base在几乎所有15个标准零样本/少样本学习任务上超越了LLaMA2 7B Base并且在大多数任务上与LLaMA3 8B Base表现相当。指令跟随能力SFT显著增强了LLaDA的指令跟随能力如多轮对话案例所示。反转推理能力LLaDA有效地打破了反转诅咒在正向和反转任务上表现一致。特别是在反转诗歌完成任务中LLaDA超越了GPT-4o。Reference[1] Large Language Diffusion Models