Training Optimal Large Diffusion Language Models (Quokka Scaling Law)¶
论文链接: https://arxiv.org/abs/2510.03280
代码/项目: https://github.com/ (官方开源基座基于 Megatron-LM 定制)
作者: Jinjie Ni, Qian Liu, Chao Du, Longxu Dou, Hang Yan, Zili Wang, Tianyu Pang, Michael Qizhe Shieh
机构: National University of Singapore (NUS), Sea AI Lab, Shanghai Qiji Zhifeng, StepFun
发表会议: NeurIPS 2026
论文分类: 大模型评测与缩放定律 / 扩散语言模型 (llm_evaluation / diffusion_models)
核心亮点 / 论文摘要¶
在自回归模型（AR）主导的大语言模型时代，扩散语言模型（Diffusion Language Models, DLMs）凭借双向上下文建模、任意阶依赖捕捉、非因果反向推理以及极具潜力的并行多 Token 解码速度，展现出追赶甚至超越 AR 架构的潜力。然而，学术界与工业界对于如何从零预训练大尺度 DLM 的基础认知几乎仍是一片空白，缺少如同 Chinchilla 之于 AR 模型的系统性缩放定律（Scaling Law）。
本论文提出了 Quokka——首个面向扩散语言模型的大规模缩放定律体系，全面覆盖计算受限（Compute-Constrained）与数据受限（Data-Constrained）两大核心机制，并深入消融了核心建模结构与优化超参： 1. 计算受限缩放定律：在固定训练计算量 \(C\) 下，最优参数量 \(N_{\text{opt}} \propto C^{0.5}\) 与数据量 \(D_{\text{opt}} \propto C^{0.5}\) 呈同速等比例扩展，与 AR 规律一致；但 DLM 的数据饥渴度是 AR 模型的 2~5 倍，在相同 FLOPs 下必须分配更多 Token 并采用相对较小的模型规模。 2. 数据受限缩放定律：首次构建并验证了跨 23,145 次实验的 U 型验证损失曲面，量化了多次 Epoch 训练引发的过拟合起始拐点，发现抗过拟合极限轮数满足 \(e_{\text{opt}} \propto U_D^{0.39} / N^{0.55}\)；例如 10B 模型在 1T 独有数据上可容忍约 1,098 轮训练而不发生泛化恶化。 3. 关键建模与优化法则：吸收型掩码转移（Absorbing Mask）在预训练收敛与下游评测（MMLU / HellaSwag）上全面碾压均匀噪声转移（Uniform Diffusion）；线性时间调度（Linear Schedule）最稳健；有原则的扩散 ELBO 变分界相比 MaskGIT 经验损失具备更优的最终泛化能力；且 AR 模型的批大小与学习率缩放法则可无缝迁移至 DLM。
背景与动机¶
近年来，随着 LLaDA 等模型的探索，离散扩散语言模型正式进入大模型缩放视野。DLM 具备固有优势：全双向注意力支持在生成过程中动态修改上下文，极大契合代码修正与多步数学推理；多 Token 并行去噪具备极低推理延迟潜力；在有限独有数据下展现出更优的数据复用能力（Data Learning Potential）。
然而，前人训练大尺度 DLM 多依赖经验直觉或生搬硬套 AR 结论。当前模型训练普遍面临两大根本问题： - 计算瓶颈场景：在固定算力预算 \(C \approx 6ND\) 下，究竟该把算力分配给模型参数量 \(N\) 还是训练 Token 数量 \(D\)？现存唯一的开源大扩散模型 LLaDA（1.1e23 FLOPs 训练 8B 模型配 2.3T 数据）配置是否达到最优点？ - 数据瓶颈场景：随着高质量网页数据线性增长速度远落后于模型算力指数增长（Token Crisis），重复多轮训练（Multi-epoch Training）成为必选项。经典 Muennighoff 等人针对 AR 拟合的数据受限公式存在致命缺陷（其假定验证损失随轮数单调不增，忽略了偏差-方差权衡下的过拟合）。对于 DLM，给定独有数据量 \(U_D\) 和模型尺寸 \(N\)，究竟最多可以复用训练多少轮？最优的 \((N, e)\) 联合配置是什么？
针对上述盲区，Quokka 通过从 7M 到 11B 参数、数据跨度 1B 到 260B Tokens 的数万次系统性实验，确立了 DLM 预训练的科学基准。
方法详解¶
整体框架¶
Quokka 的研究体系涵盖“离散扩散前向-反向过程定义”、“双重缩放定律建模（计算受限与数据受限）”以及“核心架构/超参消融”三大模块。其整体技术逻辑与最优训练配方如下流程图所示：
flowchart TD
A[大尺度离散扩散语言模型 DLM 训练需求] --> B{约束范式划分}
B -->|固定算力 C, 数据充足| C[计算受限缩放定律: IsoFLOP & 参数化拟合]
C --> C1[发现: N_opt ∝ C^0.5, D_opt ∝ C^0.5]
C1 --> C2[核心结论: DLM 比 AR 饥渴 2~5x 数据, 倾向较小模型配更大语料]
B -->|独有数据 U_D 受限, 算力充裕| D[数据受限缩放定律: U 型过拟合建模]
D --> D1[有效数据规模 D' = U_D · e^pe · exp(-...)]
D1 --> D2[过拟合拐点律: e_opt ∝ U_D^0.39 / N^0.55]
D2 --> D3[确定最优模型容量 N 与复用轮数 e 的联合配比]
A --> E[核心架构与优化策略决策]
E --> E1[转移核: 吸收型掩码 Absorbing-Mask >> 均匀噪声 Uniform]
E --> E2[噪声调度: 线性调度 α_t = 1-t 最优; 易到难课程加速收敛]
E --> E3[损失函数: 变分扩散 ELBO 最终泛化优于 MaskGIT]
E --> E4[超参迁移: 批大小与学习率法则直接沿用 AR 经验; 重复数据保留 Weight Decay]
C2 --> F[最优大扩散语言模型预训练实施指南]
D3 --> F
E1 --> F
E2 --> F
E3 --> F
E4 --> F关键设计¶
1. 吸收型离散掩码扩散前向与反向动力学
模型采用词表状态空间中的吸收掩码核（Absorbing Transition Kernel），规避了传统连续扩散或离散均匀转移难以判别真实语义的缺陷。对于长度为 \(L\) 的清晰序列 \(x_0 \in \mathcal{V}^L\)，在连续时间噪声步长 \(t \in [0, 1]\) 下引入单调递减调度函数 \(\alpha_t \in [0, 1]\)（满足 \(\alpha_0 = 1, \alpha_1 = 0\)）。前向加噪过程独立作用于每个位置：
$\(q_{t|0}(x_t^{(i)} \mid x_0^{(i)}) = \begin{cases} \alpha_t, & x_t^{(i)} = x_0^{(i)}, \\ 1 - \alpha_t, & x_t^{(i)} = [\text{MASK}]. \end{cases}\)$
反向去噪过程从全掩码状态 \(x_1\) 逐步推进至 \(x_0\)。该动力学具备时间无关特性（Time-Agnostic Property），即未被掩码的上下文决定了真实 Token 的条件概率，去噪器无需复杂的连续时间显式位置嵌入，天然支持并行双向自注意力架构。
2. 计算受限下的幂律双路径协同标定
在固定总 FLOPs \(C \approx 6ND\) 下，作者通过两种独立方法相互印证：
- 方法一（IsoFLOPs 轮廓线）：设定 9 档固定 FLOPs 预算（\(3 \times 10^{18}\) 至 \(1 \times 10^{21}\)），在对数坐标系下密集扫描模型大小 \(N\)，对每条 IsoFLOP 曲线拟合抛物线提取极小值点，进而外推幂律关系 \(N_{\text{opt}} = k_N C^a, D_{\text{opt}} = k_D C^b\)。
- 方法二（参数化损失曲面拟合）：基于经典风险分解理论拟合双因子损失函数：
$\(L(N, D) \triangleq E + \frac{A}{N^\alpha} + \frac{B}{D^\beta}\)$
其中 \(E\) 为不可约误差（Irreducible Error，反映数据固有熵与扩散变分下界固有差距），\(A/N^\alpha\) 代表受限参数容量引起的模型误差，\(B/D^\beta\) 代表有限 Token 采样引起的训练误差。通过 L-BFGS 算法最小化 Huber 损失，得到解析极值解：
$\(a = \frac{\beta}{\alpha+\beta}, \quad b = \frac{\alpha}{\alpha+\beta}, \quad G = \left(\frac{\alpha A}{\beta B}\right)^{\frac{1}{\alpha+\beta}}, \quad N_{\text{opt}}(C) = G \left(\frac{C}{6}\right)^a, \quad D_{\text{opt}}(C) = G^{-1} \left(\frac{C}{6}\right)^b\)$
3. 数据受限与 U 型过拟合泛化曲面解析
针对有限独有数据 \(U_D\) 和多次训练 Epoch \(e\)，经典模型无法预测过拟合回升。作者引入结合“学习增益项”与“过拟合惩罚项”的有效数据规模 \(D'\)，构建全周期模型：
$\(L(N, U_D, e) \triangleq E + \frac{A}{N^\alpha} + \frac{B}{D'^\beta}\)$
$\(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}}\)$
在训练初期（\(e\) 较小），学习项 \(e^{p_e}\) 占主导，有效数据量增长促使验证损失下降；当轮数超过峰值过拟合周期 \(e_p\) 时，指数衰减项占据主导，有效数据规模萎缩，损失曲线反弹呈现典型的 U 型曲线。该设计在单轮训练（\(e \le 1\)）时自然退化为标准计算受限定律。
一个完整示例¶
考虑在工业级场景下分配算力：假设拥有与开源模型 LLaDA 相同的算力预算 \(C = 1.1 \times 10^{23}\) FLOPs。 按照传统经验或直接采用 8B 参数模型，LLaDA 在 2.3T Tokens 上进行训练。然而，将 \(C = 1.1 \times 10^{23}\) 代入 Quokka 计算受限法则： - 由方法一公式 \(N_{\text{opt}} \approx 0.0216 \times C^{0.514}\)，计算可得算力最优参数量为 \(N_{\text{opt}} \approx 15\text{B}\)。 - 代入 \(D_{\text{opt}} \approx 7.7 \times C^{0.486}\)，最优训练数据量约为 \(D_{\text{opt}} \approx 1.2\text{T}\) Tokens。 这表明 LLaDA 实际上属于典型的“过度训练小模型”配置（模型较最优配置偏小近 2 倍，语料偏多近 2 倍），在给定总算力下并未达到预训练损失的理论最优点。
损失函数 / 训练策略¶
在扩散模型训练中，去噪器参数化为条件概率分布 \(p_\theta(x_0^{(i)} \mid x_t)\)。变分证据下界（ELBO）导出的严格扩散损失表示为重要性加权形式： $\(\mathcal{L}_{\text{diff}} = \int_0^1 w(t; \alpha) \mathbb{E}_{q_{t|0}(x_t \mid x_0)} \left[ \sum_{i: x_t^{(i)} = [\text{MASK}]} -\log p_\theta(x_0^{(i)} \mid x_t) \right] \mathrm{d}t\)$ 其中权重因子 \(w(t; \alpha) = \frac{\alpha'_t}{\alpha_t - 1}\) 用于补偿不同加噪时间点期望掩码位置数量的不均衡。在线性调度 \(\alpha_t = 1 - t\) 下，权重自然简化为 \(w(t) = \frac{1}{t}\)。
训练策略关键细节包括： - 课程加噪策略（Easy-to-Hard Curriculum）：在训练初期利用动态滑动的单峰高斯窗口优先对小 \(t\)（少掩码、低噪声）进行密集采样，使网络先学习短距强关联填空，随后平滑过渡至全范围均匀采样，显著提升训练初期的收敛速率并在末期斩获微弱的下游精度提升。 - 权重衰减与数值稳定性：在单 Epoch 预训练中，去掉 Weight Decay 实际上收敛更快；但在数据受限的多轮重复训练（如 96 轮）中，Weight Decay 成为抑制模型范数暴增的定海神针。尤其在 BF16 精度下，由于其尾数位仅 7 位，参数模长过大导致 Softmax 前 Logits 超过 128 时会发生严重的量化截断与粗化，故高频数据复用训练必须保留 Weight Decay。
实验结果与分析¶
论文在基于 Megatron-LM 深度重构的架构下开展了详尽的实验，评测指标涵盖验证集交叉熵损失（C4-val / Nemotron-CC）以及主流下游零样本/少样本基准 MMLU 和 HellaSwag。
1. 计算受限缩放系数对比¶
表 1 详细展示了 Quokka 与经典的自回归 Chinchilla 定律在各项幂律参数上的量化对比。可以看出，DLM 与 AR 模型的指数 \(a, b\) 高度接近 0.5，但 DLM 的数据比例系数 \(k_D\) 高出 Chinchilla 2.2~6.7 倍。
| 模型与估计方法 | 不可约损失 \(E\) | 模型指数 \(a\) | 数据指数 \(b\) | 系数 \(k_N\) | 系数 \(k_D\) |
|---|---|---|---|---|---|
| Chinchilla Approach 1 | - | 0.50 | 0.50 | 0.09 | 1.88 |
| Chinchilla Approach 2 | - | 0.49 | 0.51 | 0.15 | 1.15 |
| Chinchilla Approach 3 (Parametric) | 1.69 | 0.46 | 0.54 | 0.60 | 0.28 |
| Quokka Approach 1 (IsoFLOPs) | - | 0.51 | 0.49 | 0.02 | 7.70 |
| Quokka Approach 2 (Parametric) | 2.41 | 0.50 | 0.50 | 0.04 | 4.10 |
注：Quokka 参数化拟合损失公式为 \(L(N, D) \approx 2.413 + \frac{798.6}{N^{0.379}} + \frac{4604.9}{D^{0.378}}\)。不可约损失 \(E=2.41\) 高于 AR 的 1.69，反映了离散扩散变分界与加噪离散化带来的固有理论间隙。
2. 计算最优预算与 Token 分配对比¶
表 2 给出了从 400M 到 10T 参数规模下，Quokka（基于方法一预测）与 Chinchilla 的对标数据分配推荐。对于相同的目标参数量，DLM 达成算力最优所需的 Token 数与训练算力普遍达到 AR 模型的 3~5 倍。
| 参数量规模 (Parameters) | Quokka 最优算力 (FLOPs) | Quokka 最优数据 (Tokens) | Chinchilla 算力 (FLOPs) | Chinchilla 数据 (Tokens) |
|---|---|---|---|---|
| 400 M | \(9.46 \times 10^{19}\) | 39.3 B | \(1.92 \times 10^{19}\) | 8.0 B |
| 1 B | \(5.62 \times 10^{20}\) | 93.5 B | \(1.21 \times 10^{20}\) | 20.2 B |
| 10 B | \(4.96 \times 10^{22}\) | 825.2 B | \(1.23 \times 10^{22}\) | 205.1 B |
| 67 B | \(2.01 \times 10^{24}\) | 5.0 T | \(5.76 \times 10^{23}\) | 1.5 T |
| 175 B | \(1.30 \times 10^{25}\) | 12.4 T | \(3.85 \times 10^{24}\) | 3.7 T |
| 280 B | \(3.24 \times 10^{25}\) | 19.3 T | \(9.90 \times 10^{24}\) | 5.9 T |
| 520 B | \(1.08 \times 10^{26}\) | 34.6 T | \(3.43 \times 10^{25}\) | 11.0 T |
| 1 T | \(3.86 \times 10^{26}\) | 64.2 T | \(1.27 \times 10^{26}\) | 21.2 T |
| 10 T | \(3.41 \times 10^{28}\) | 566.4 T | \(1.30 \times 10^{28}\) | 216.2 T |
3. 数据受限场景下的耐受轮数与联合参数分配¶
在数据受限拟合方程（23,145 次实验拟合结果，Huber 损失 9.78 优于替代形式的 23.8 与 31.52）中： $\(L(N, U_D, e) = \frac{1535.23}{N^{0.42}} + \frac{54.21}{\left( U_D \cdot e^{1.49} \cdot \exp\left( - \left( \frac{\max(0, e-1)}{254.35 \frac{U_D^{0.39}}{N^{0.55}}} \right)^{0.40} \right) \right)^{0.13}}\)$
表 3 明确解答了“给定模型尺寸与独有数据预算，最多可训练多少轮而泛化不恶化”：
| 模型参数 \(N\) \ 独有数据 \(U_D\) | 10 M | 100 M | 1 B | 10 B | 100 B | 1 T | 10 T |
|---|---|---|---|---|---|---|---|
| 400 M | 70 | 175 | 430 | 1,057 | 2,593 | 6,357 | 15,585 |
| 1 B | 42 | 105 | 260 | 641 | 1,572 | 3,857 | 9,456 |
| 10 B | 11 | 29 | 73 | 181 | 447 | 1,098 | 2,693 |
| 67 B | 1 | 9 | 25 | 63 | 157 | 388 | 953 |
| 175 B | 1 | 4 | 14 | 37 | 93 | 229 | 564 |
| 1 T | 1 | 1 | 4 | 13 | 35 | 88 | 217 |
观察发现：独有数据量越大，过拟合延迟出现；模型参数量越大，过拟合加速显现。
在独有数据 \(U_D\) 固定且算力无上限的前提下，表 4 进一步给出了参数量 \(N\) 与训练轮数 \(e\) 的全局联合最优配比：
| 独有 Token 预算 \(U_D\) | 最优参数量 \(N_{\text{opt}}\) | 最优训练轮数 \(e_{\text{opt}}\) | 等效消耗算力 FLOPs |
|---|---|---|---|
| 10 M | 41 M | 247 | \(6.07 \times 10^{17}\) |
| 100 M | 95 M | 356 | \(2.04 \times 10^{19}\) |
| 1 B | 222 M | 569 | \(7.59 \times 10^{20}\) |
| 10 B | 518 M | 910 | \(2.83 \times 10^{22}\) |
| 100 B | 1.6 B | 1,151 | \(1.11 \times 10^{24}\) |
| 1 T | 3.7 B | 1,842 | \(4.12 \times 10^{25}\) |
| 10 T | 8.7 B | 2,947 | \(1.54 \times 10^{27}\) |
局限性与思考 / 讨论¶
- 变分损失与自回归交叉熵的绝对值差异：DLM 拟合得到的不可约损失 \(E \approx 2.41\) 显著高于 AR 模型的 1.69。这是因为离散扩散以 ELBO 作为代理目标，加噪扰动与去噪采样引入了变分下界松弛（Variational Gap），导致绝对损失值无法与 AR 直接对标，评测应更多关注相对收敛趋势与下游基准。
- 多轮训练的偶发双峰下降（Double Descent）：在极少数超长轮数训练中，数据受限实验观察到了二次下降现象。本文公式主要捕捉经典的一阶 U 型拐点，对于复杂双峰动力学的理论解析仍待后续工作补充。
- 推理延迟与训练代价的工业权衡：尽管缩放定律推荐使用更小的模型参数和更多的训练数据（对 DLM 而言 2~5 倍数据意味着更长的前期训练开销），但在实际部署中，较小的模型不仅显存占用更低，而且配合 DLM 的并行解码与 KV Cache 优化机制，能带来显著的端到端吞吐收益。
总结¶
Quokka 为扩散语言模型构建了首个坚实的缩放理论与工程实践坐标系： - 确立了计算受限下的 \(N \propto C^{0.5}, D \propto C^{0.5}\) 演进轨迹，明确了 DLM “偏好较小模型、投喂海量语料”的 2~5 倍数据饥渴特性； - 解析了数据受限下的过拟合动力学，以 \(e_{\text{opt}} \propto U_D^{0.39} / N^{0.55}\) 为 Token 危机下的数据复用提供了定量安全边界； - 为吸收型掩码设计、线性调度器、变分 ELBO 损失以及权重衰减的配置提供了详实的实证支撑。这一工作不仅填补了 DLM 大规模预训练的经验法则空白，也为未来非自回归与扩散大语言模型的发展提供了不可或缺的底层指南。