Training Optimal Large Diffusion Language Models¶

会议: NeurIPS 2026
arXiv: 2510.03280
OpenReview: https://openreview.net/forum?id=VV80K0h5rJ
领域: LLM 其他
关键词: 扩散语言模型, 缩放定律, 计算最优, 数据受限, 吸收掩码

一句话总结¶

提出了首个涵盖计算受限与数据受限两大体制的离散扩散语言模型大规模缩放定律 Quokka，指出扩散模型相比自回归模型更为数据饥渴（需分配 2–5 倍训练数据），并揭示了过拟合 U 型验证损失曲面及最优建模选择。

研究背景与动机¶

随着离散扩散模型在文本生成领域的突破，以 LLaDA、Gemini Diffusion 和 Mercury 为代表的扩散语言模型（Diffusion Language Models, DLMs）展示出巨大的应用潜力。DLMs 依托双向注意力与任意顺序生成机制，在代码补全、数学推理和低延迟并行解码等非因果依赖场景中展现出优于自回归（Autoregressive, AR）模型的特性。特别是在多轮数据复用与长尾语料学习上，扩散机制展现出独特的抗过拟合潜质。然而，当前业界对如何从零开始高效训练数十亿甚至数百亿参数的超大规模 DLMs 仍缺乏系统性理论指导与经验法则，既有尝试大多停留在启发式调优或机械照搬自回归模型的超参数与缩放配比。

这种认知盲区的核心矛盾在于：大语言模型的训练预算本质上由算力和语料双重塑造。在固定计算预算下，从业者亟需知道模型参数量与训练 Token 数的帕累托最优分配；而在现实面临的优质文本语料枯竭（Token Crisis）困境下，又必须厘清重复多轮训练何时触碰过拟合拐点，以及有限唯一语料下参数量与训练 Epoch 的联合最优配置。既有针对自回归模型的 Chinchilla 定律假设语料无限且单轮遍历，而 Muennighoff 等人的数据受限扩展公式单调递减、无法刻画过拟合上升，均不能直接套用于具有破坏-去噪特性的离散扩散模型。

为了填补这一理论与工程空白，本文系统构建了面向扩散语言模型的首个全尺度缩放定律体系 Quokka。作者跨越从 7M 到 11B 的模型尺寸与高达 260B 的 Token 规模，基于两万余次实验全景描摹了计算与数据约束下的损失曲面。核心 idea：建立首个覆盖计算最优与数据受限体制的扩散语言模型缩放定律 Quokka，证明在相同算力下 DLMs 需配置 2–5 倍于 AR 模型的数据量以实现计算最优，并提出可精准捕捉过拟合拐点的 U 型有效数据量验证损失模型。

方法详解¶

整体框架¶

Quokka 从计算受限（Compute-Constrained）与数据受限（Data-Constrained）两个宏观维度，配合微观层面的核心建模与优化超参消融，全面确立了训练最优扩散语言模型的规范体系。其整体评估与推导流水线包含四个核心步骤：首先通过单轮遍历下的 IsoFLOPs 轮廓线与参数化损失拟合确立计算最优边界；其次通过多轮重复训练刻画泛化误差与过拟合动态，建立带峰值 Epoch 惩罚项的 U 型损失曲面；再次横向对比转移核类型与时间步调度策略；最后验证优化超参数（Batch Size、学习率、权重衰减）从自回归模型向扩散模型的迁移规律。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["训练预算输入<br/>固定总 FLOPs C 或唯一语料 U_D"] --> B["计算受限缩放分析<br/>IsoFLOPs 极值点 + 参数化拟合"]
    B --> C["数据受限 U 型建模<br/>学习项递增 + 过拟合指数惩罚"]
    C --> D["建模设计消融<br/>吸收掩码核 vs 均匀核 / 线性调度"]
    D --> E["超参迁移律与正则化<br/>AR 批大小与学习率迁移 / 权重衰减控模"]
    E --> F["训练最优 DLM 配置输出<br/>最优参数量 N_opt、Token 预算 D_opt、Epoch e_opt"]

关键设计¶

1. 计算受限缩放定律：双途径推导与数据饥渴性建模 针对固定浮点运算预算 \(C\) 下模型参数量 \(N\) 与训练数据量 \(D\) 如何权衡的经典问题，Quokka 采用 IsoFLOPs 轮廓剖面与参数化风险分解双重途径进行交叉验证。在 IsoFLOPs 实验中，作者固定 9 组算力预算（从 \(3\times 10^{18}\) 至 \(1\times 10^{21}\) FLOPs），通过拟合抛物线极小值捕捉最优点；在参数化拟合中，将经验损失分解为不可约误差 \(E\)、模型容量误差 \(A/N^\alpha\) 与数据优化误差 \(B/D^\beta\)：

\[L(N, D) \approx 2.413 + \frac{798.6}{N^{0.379}} + \frac{4604.9}{D^{0.378}}\]

双途径一致得出缩放指数 \(a \approx 0.50\) 与 \(b \approx 0.50\)（即 \(N_{\text{opt}} \propto C^{0.5}\) 且 \(D_{\text{opt}} \propto C^{0.5}\)），说明模型规模与数据量应随算力同速率等比扩展。但与 Chinchilla 定律的关键差异在于比例系数：DLMs 的数据乘子 \(k_D\) 高达 \(4.10 \sim 7.70\)，为 Chinchilla（\(1.15 \sim 1.88\)）的 2.2 至 6.7 倍。这意味着在相同单轮训练算力下，扩散模型因输入人为加噪破坏了部分信息，表现出更强烈的数据饥渴性，最优策略必须选用参数量更小、训练语料显著更大的组合。

2. 数据受限与过拟合 U 型损失曲面：有效数据规模与峰值 epoch 建模 面对真实世界中唯一语料 \(U_D\) 有限而算力充裕的场景，重复多轮训练必然面临偏差-方差权衡诱发的过拟合。既有自回归数据受限公式假设损失单调非递增，无法反映性能劣化。Quokka 基于 23,145 次大规模训练轨迹，首创将有效数据规模 \(D'\) 解耦为「有效学习乘子」与「过拟合惩罚指数」的乘积：

\[D' = U_D \cdot e^{p_e} \cdot \exp\left(-\left(\frac{\max(0, e-1)}{e_p}\right)^\gamma\right), \quad e_p = c_p \frac{U_D^{m_p}}{N^{k_p}}\]

其中峰值拐点 Epoch \(e_p\) 由唯一语料量与模型尺寸共同决定。拟合结果显示 \(e_{\text{opt}} \propto U_D^{0.39} / N^{0.55}\)，说明增加高质量语料能显著推迟过拟合发生，而扩充模型容量则加速过拟合。该公式在 \(e \le 1\) 时指数惩罚项严格退化为 1，无缝回归计算受限规律。借助该闭式解，能够精确预估任意规模模型在特定语料池下的最大可承受训练轮数（例如 10B 模型在 1T 语料上可支持约 1098 轮有效训练而不退化）。

3. 吸收掩码状态转移核：优于均匀离散扩散的高效去噪范式 在离散状态转移矩阵的设计上，论文在 1B 模型规模与 96B 语料上严格对比了「均匀转移核（Uniform Kernel）」与「吸收掩码转移核（Absorbing Mask Kernel）」。均匀核通过全词表均匀随机替换将文本破坏为白噪声，去噪过程需学习将任意错误嵌入映射回真实词；而吸收掩码核仅允许有效词向固定 \([MASK]\) 标记单向转移，去噪等价于完形填空。实验表明，掩码扩散在预训练收敛损失与下游推理基准上均全面大幅碾压均匀扩散。对均匀扩散细粒度分析发现，网络虽然能轻易分辨哪些位置被噪声替换，却难以将离散随机词向量重构为真实目标，从而确立了吸收掩码在大规模文本扩散建模中的基石地位。

4. 扩散调度、渐进曲率与损失目标：线性调度与 ELBO 变分重加权 针对正向加噪速率曲线 \(\alpha_t\)，对比了线性（\(\alpha_t = 1-t\)）、二次多项式（\(\text{poly2}, \alpha_t = 1-t^2\)）与余弦（\(\text{cosine}\)）三种轨迹。线性调度由于具有平衡的加噪斜率和稳定的梯度方差，在综合基准上取得最优且最具稳定性的表现；poly2 调度在常识推理上微弱占优，而过度偏向重噪声的余弦调度表现最差。此外，引入基于滑动高斯窗口的「由易到难」时间步采样课程（早期偏向低噪声、后期偏向高噪声）可加速初期收敛并提升最终性能。在优化目标层面，尽管不带时间步重要性加权的 MaskGIT 损失在训练初期凭借更均匀的梯度下降表现出快速收敛，但基于严格证据下界（ELBO）重要性加权 \(w(t;\alpha) = \frac{\alpha'_t}{\alpha_t - 1}\) 的变分扩散损失在长序列充分训练后最终达成更低的上界与更强性能。

损失函数 / 训练策略¶

离散掩码扩散模型的连续时间学习目标通过最大化对数似然的变分证据下界（ELBO）导出，离散化后转化为加权交叉熵：

\[\mathcal{L} = \int_{0}^{1} w(t;\alpha)\; \mathbb{E}_{q_{t|0}(x_t \mid x_0)}\left[\sum_{i: x_t^{(i)}=m} -\log p_\theta\left(x_0^{(i)} \mid x_t\right)\right] \mathrm{d}t\]

对于标准线性调度 \(\alpha_t = 1-t\)，其天然重要性权重精确简化为 \(w(t) = 1/t\)。

在优化超参数方面，作者证实自回归模型建立的超参规律能够无缝迁移至 DLMs： - 批大小（Batch Size）：1B 模型在 256 与 1024 批大小下表现相近，而 4096 批大小收敛显著落后，表明扩散目标并未改变临界批大小特性； - 学习率（Learning Rate）：峰值学习率设为 \(1\times 10^{-4} \sim 2\times 10^{-4}\) 配合余弦退火具有最高通用性； - 权重衰减（Weight Decay）：在单轮训练中关闭权重衰减甚至略微加速收敛；但在多轮数据复用训练中，开启权重衰减对于抑制参数 L2 范数爆炸、防止 bf16 浮点格式在 Logits 极大时发生数值下溢与坍塌具有不可替代的稳定性保障。

实验关键数据¶

主实验¶

Quokka 缩放定律推导出的系数与 Chinchilla 的对比，以及各参数规模下对应的计算最优资源分配如下表所示。

评估维度 / 设定	拟合方法	不可约损失 \(E\)	参数指数 \(a\)	数据指数 \(b\)	乘子 \(k_N\)	乘子 \(k_D\)
Chinchilla (AR)	IsoFLOPs / 双变量	-	0.50	0.50	0.09	1.88
Chinchilla (AR)	参数化损失拟合	1.69	0.46	0.54	0.60	0.28
Quokka (DLM)	IsoFLOPs 轮廓剖面	-	0.51	0.49	0.02	7.70
Quokka (DLM)	参数化损失拟合	2.41	0.50	0.50	0.04	4.10

在具体模型规模下的计算最优分配对比如下表：

模型参数量 \(N\)	Quokka 最优算力 (FLOPs)	Quokka 最优 Token 数 \(D\)	Chinchilla 算力 (FLOPs)	Chinchilla Token 数 \(D\)
400M	\(9.46 \times 10^{19}\)	39.3 B	\(1.92 \times 10^{19}\)	8.0 B
1B	\(5.62 \times 10^{20}\)	93.5 B	\(1.21 \times 10^{20}\)	20.2 B
10B	\(4.96 \times 10^{22}\)	825.2 B	\(1.23 \times 10^{22}\)	205.1 B
67B	\(2.01 \times 10^{24}\)	5.0 T	\(5.76 \times 10^{23}\)	1.5 T
175B	\(1.30 \times 10^{25}\)	12.4 T	\(3.85 \times 10^{24}\)	3.7 T

消融实验¶

在数据受限设定下，模型规模与唯一语料约束对可承受的最大训练 Epoch \(e_{\text{opt}}\)（过拟合拐点）的预测如下表所示：

参数量 \(N\) \ 唯一语料 \(U_D\)	100M Tokens	1B Tokens	10B Tokens	100B Tokens	1T Tokens
400M	175 epochs	430 epochs	1057 epochs	2593 epochs	6357 epochs
1B	105 epochs	260 epochs	641 epochs	1572 epochs	3857 epochs
10B	29 epochs	73 epochs	181 epochs	447 epochs	1098 epochs
67B	9 epochs	25 epochs	63 epochs	157 epochs	388 epochs
175B	4 epochs	14 epochs	37 epochs	93 epochs	229 epochs

在固定唯一语料预算 \(U_D\)、算力不受限条件下的联合最优配置预测：

唯一语料规模 \(U_D\)	最优参数量 \(N_{\text{opt}}\)	最优训练轮数 \(e_{\text{opt}}\)	对应训练算力 (FLOPs)
100M Tokens	95 M	356 epochs	\(2.04 \times 10^{19}\)
1B Tokens	222 M	569 epochs	\(7.59 \times 10^{20}\)
10B Tokens	518 M	910 epochs	\(2.83 \times 10^{22}\)
100B Tokens	1.6 B	1151 epochs	\(1.11 \times 10^{24}\)
1T Tokens	3.7 B	1842 epochs	\(4.12 \times 10^{25}\)

关键发现¶

数据饥渴特性明确：在固定 FLOPs 预算下，DLMs 达成最优性能需要分配 \(2 \sim 5\) 倍于 AR 模型的数据量。例如在当前开源基准 LLaDA 的算力预算（\(1.1\times 10^{23}\) FLOPs）下，其采用的 8B 参数与 2.3T 数据属于欠训练与过大模型配置，Quokka 测算出的最优配比应为 15B 参数模型搭配 1.2T 语料。
不可约损失显著高于 AR：Quokka 拟合出的扩散模型不可约损失 \(E \approx 2.41\)，高于 Chinchilla 的 1.69。这是因为扩散模型优化的是包含加噪变分隙的 ELBO 上界，而非自回归的精确负对数似然。
过拟合拐点对语料极度敏感：过拟合起始轮数服从 \(e_{\text{opt}} \propto U_D^{0.39}/N^{0.55}\)。扩大唯一语料规模是延缓过拟合的最有效手段，而大参数模型在有限语料上重复训练极易提前陷入过拟合。

亮点与洞察¶

填补扩散语言模型规模化理论空白：给出了学术界与工业界训练超大离散扩散语言模型的第一份定量路线图，终结了依靠直觉或机械搬运自回归经验的粗放尝试。
揭示了加噪去噪的信息重构代价：从第一性原理阐明了 DLM 数据饥渴的根源在于前向破坏过程引入了有效信息衰减，必须借助更高的语料规模来对冲掩码带来的有效信号稀释。
开创性提出过拟合 U 型损失曲面公式：巧妙通过解耦学习增长项与指数衰减惩罚项，不仅完美刻画了多轮训练下的过拟合拐点，还能在单轮时退化为标准 Chinchilla 定律。

局限与展望¶

验证集与分词器依赖度高：拟合参数绝对数值受验证集分布及 GPT-2 分词器特性影响较大，在跨领域或更换新型 Byte 分词器时可能发生数值偏移。
局部存在二次下降（Double Descent）现象：极少数长周期超多轮次运行中观察到双重下降趋势，拟合公式目前仅截断建模了单峰上升拐点，理论完整性仍有提升空间。
缺少强化学习后训练维度的缩放律：目前探索集中于预训练阶段，面向思维链推理或 Test-Time Diffusion 的推理期计算缩放仍待后续深入。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次系统性推导扩散语言模型在计算和数据双重约束下的全景缩放定律。
实验充分度: ⭐⭐⭐⭐⭐ 跨越从 7M 到 11B 参数量，包含两万三千余次独立训练运行，证据极为扎实。
写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严整，论证层层递进，理论公式与工程指引高度统一。
价值: ⭐⭐⭐⭐⭐ 为未来离散扩散语言模型的架构选型、算力预估与预训练规划提供了不可或缺的基准指南。