Looped Diffusion Language Models¶

会议: NeurIPS 2026
arXiv: 2605.26106
OpenReview: https://openreview.net/forum?id=UqgMsw7AZH
论文: https://arxiv.org/abs/2605.26106
领域: 预训练
关键词: 掩码扩散语言模型, 循环 Transformer, 深度缩放, 掩码工作区, 自适应测试时计算

一句话总结¶

本文提出了 LoopMDM，首次将循环机制引入掩码扩散语言模型（MDM），通过仅对少量前中部 Transformer 层进行权重共享的循环计算，在相同训练 FLOPs 下大幅提升预训练效率与推理质量，并利用掩码位置作为并行计算工作区实现了无需增加参数的测试时计算缩放。

研究背景与动机¶

掩码扩散语言模型（Masked Diffusion Models, MDMs）近期作为自回归模型（ARMs）的有力替代者迅速兴起，并在 LLaDA、Dream 等大模型上展现出媲美 ARMs 的建模能力。现有的优化工作主要集中在去噪训练目标、掩码噪声调度（noise schedule）以及更深更宽的架构扩展上。然而，这种单纯依赖模型参数量与训练 token 量的扩展范式，往往忽视了扩散模型本身天然具备的去噪迭代计算特性，亟需从模型架构层面探索新的计算效率维度。

在自回归文献中，循环 Transformer（Looped Transformer）通过重复执行共享网络块，以恒定的参数量换取更深的等效计算深度，已被广泛用作测试时计算缩放工具。但在 MDM 场景下，循环架构的价值长期未被探索，甚至存在直觉上的质疑：掩码扩散模型本身就需要在外部去噪步（denoising steps）之间执行多轮迭代，在单步去噪网络内部再引入循环是否纯属计算冗余？更重要的是，以往的循环模型多聚焦于纯粹的参数效率或推理精炼，在同等训练 FLOPs 预算下，循环模型往往因每步计算量增加导致训练 token 曝光不足，最终性能很难击败无循环基线。

针对上述核心矛盾，本文深入研究了在严格对齐总训练 FLOPs 的前提下循环机制对 MDM 的影响。研究发现，盲目对整个网络做全盘循环并不可行，但仅对前中部（early-middle）的少量层引入选择性循环，便能以极低的额外计算代价激活深度缩放红利，并促使模型将未解析的掩码 token 视作并行隐式计算的“工作区”（workspace）。核心 idea：本文提出了首个面向掩码扩散语言模型的循环架构 LoopMDM，通过仅将去噪网络前中部的少量层作为权重共享的循环块并在训练时随机采样循环次数，使掩码位置充当协同推理的动态工作区，在对齐训练总 FLOPs 的条件下以高达 3.3 倍的计算效率超越同规模非循环模型，并支持在推理阶段按需动态扩展循环深度。

方法详解¶

整体框架¶

LoopMDM 将掩码扩散模型的去噪网络划分为按序连接的三大组件：输入头部（Head，含 \(n_h\) 层）、循环中部（Mid-block，含 \(n_m\) 层）以及输出尾部（Tail，含 \(n_o\) 层，包含线性投影 LM Head）。其中 Head 和 Tail 在单次去噪前向传播中仅执行一次，分别保留对浅层 token 嵌入与最终词表分布预测的特化映射能力；而 Mid-block 则作为核心循环算子，在完全共享同一套参数的前提下，将隐状态反复迭代执行 \(S\) 次循环。

对于输入长度为 \(L\) 的含掩码序列 \(x_t \in \bar{\mathcal{V}}^L\)，Head 首先将其编码为初始隐表征 \(H_t^{(0)} \in \mathbb{R}^{L \times d}\)；随后 Mid-block 以自身上一轮输出为输入连续执行 \(S\) 轮循环迭代计算；最后 Tail 接收第 \(S\) 轮的最终隐状态 \(H_t^{(S)}\) 并预测所有掩码位置上的干净 token 概率分布 \(p_\theta^S(x_0 \mid x_t)\)。整个前向网络展开后的等效计算深度为 \(n_h + S \cdot n_m + n_o\)，而实际模型参数量与深度仅为 \(n_h + n_m + n_o\) 的基线完全一致。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入含掩码序列<br/>xt 在 V_bar^L"] --> B["输入特化头部 Head<br/>单次前向编码 Ht^(0)"]
    B --> C["1. 选择性前中部循环 Mid-block<br/>共享权重连续执行 S 轮迭代"]
    C -->|递归更新 Ht^(k)| C
    C --> D["输出特化尾部 Tail<br/>接收 Ht^(S) 并映射至词表"]
    D --> E["2. 随机循环训练调度<br/>S ~ U{1,...,Smax} 计算 NELBO"]
    D --> F["3. 掩码交互与自适应推理<br/>监测隐状态变化率动态早停"]
    E --> G["无损训练 FLOPs 对齐优化"]
    F --> H["最终去噪预测分布 p_theta"]

关键设计¶

1. 选择性前中部循环：在特征稳定层释放复用红利

以往循环 Transformer 常尝试对全网络层进行循环堆叠，但这会导致每步训练代价飙升，使得模型在固定总 FLOPs 预算下因训练步数过少而性能严重劣化。本文通过系统消融发现，循环收益对网络层所处的位置极为敏感：循环紧邻输入的嵌入层或紧邻输出的预测层均会破坏模型对基础 token 的特化表达与最终分类校准。

LoopMDM 确立了选择性循环准则：仅将循环限制在前中部（例如 12 层 DiT 中的第 1 至 2 层，基于 0 索引），且循环层数 \(n_m\) 维持在极小规模（如 \(n_m = 2\)）。这一设计的物理逻辑在于：浅层特征在刚刚完成底层 token 融合后进入一个表示相对稳定的表征空间，在此处进行共享算子的重复变换，能够以最低的训练增量 FLOPs 换取最大的计算深度收益，避免深层特化结构被强制循环共享而损坏。

2. 随机循环训练调度：构建跨有效深度的平滑算子

为了让 Mid-block 成为一个鲁棒的迭代精炼算子，而不是过拟合于某一个固定的展开深度，LoopMDM 在预训练期间采用了离散均匀随机采样策略。在每个训练步，从离散集合 \(\{1, \dots, S_{\max}\}\) 中独立均匀采样单步循环次数 \(S \sim \mathcal{U}\{1, \dots, S_{\max}\}\)，并仅对第 \(S\) 次循环的最终输出计算负证据下界损失（NELBO）：

\[\mathcal{L}_{\text{NELBO}} = \mathbb{E}_{t, S, x_t} \left[ \sum_{i=1}^{L} \mathbb{I}[x_t^i = \mathrm{m}] \frac{\alpha'_t}{1 - \alpha_t} \log \langle p_\theta^{S,i}(x_t, t), x_0^i \rangle \right]\]

在该随机机制下，Mid-block 单步的平均计算开销对应期望循环次数 \(\mathbb{E}[S] = (S_{\max} + 1) / 2\)。为了实现公平的等算力对比（Iso-FLOP），LoopMDM 相应地将总训练步数按比例缩减，确保总训练 FLOPs 与标准 MDM 严格持平。更重要的是，这种多深度随机暴露赋予了模型强大的跨计算预算泛化能力，使其在推理时不仅能任意指定 \(S\)，甚至能在外推至 \(S > S_{\max}\) 时继续保持性能单调增益。

3. 掩码位置作为工作区与自适应推理早停：动态分配计算预算

在理论层面，掩码扩散模型包含大量未解析的掩码 token，这些掩码 token 的连续隐状态在未被离散解掩码前充当了并行的“计算工作区”（Masked Workspace）。注意力分析证明，随着循环次数 \(S\) 增加，Query 掩码与 Key 掩码之间的自注意力权重（Mask-to-Mask Attention）显著提升并在 \(S_{\max}\) 附近饱和，表明内部循环的本质是让模型在未决定的 token 之间充分传递信息、消除局部与全局约束的冲突。

此外，时序分析表明循环计算在不同扩散时刻 \(t\) 上的收益并不均等：在去噪中间阶段（中间噪声水平，上下文信息初步具备但未决 token 仍然较多），循环带来的 NLL 增益最为陡峭，而在两端（几乎全掩码或几乎无掩码）增益较平缓。基于此，LoopMDM 提出了一种基于隐状态收敛性的轻量级自适应推理策略：在推理时，每轮循环监测隐状态的相对变化率：

\[\frac{\| H_t^{(k)} - H_t^{(k-1)} \|}{\| H_t^{(k)} \|}\]

当该相对改变量低于设定的阈值 \(\epsilon\)（论文推荐 \(\epsilon = 0.1\)）或达到预设的最大循环预算时，提前终止当前去噪步的循环并输出隐状态。这一机制在无需外部时序先验的情况下，自动将算力倾斜分配给最需要深度推理的中间去噪步，在几乎无损下流水准的前提下将平均推理循环步数削减一半以上。

损失函数 / 训练策略¶

模型采用基于连续时间 \(t \in [0, 1]\) 的标准掩码扩散损失目标，前向加噪过程采用单调递减调度函数 \(\alpha_t\)，在掩码 token 处根据真实 token 的对数似然加权更新。优化器选用 AdamW，预训练学习率预热至 \(3 \times 10^{-4}\)，批大小为 512，搭配序列拼接（sequence packing）技术。在总训练预算控制上，标准非循环基线在 LM1B、OpenWebText 和 FineWeb-Edu 上训练 1M 步，而带有 \(n_m=2, S_{\max}=12\) 的 LoopMDM 则按照公式 \(F_{\text{loop}} = F_{\text{base}} + (\mathbb{E}[S] - 1) n_m F_{\text{layer}}\) 扣减训练步数，使两者的总浮点运算次数严格恒等。

实验关键数据¶

主实验¶

主实验在 FineWeb-Edu 预训练模型（170M 参数）上评估了不同循环次数 \(S\) 下的零样本困惑度（Zero-shot PPL）及 9 个通用下游基准任务表现。所有模型均在严格相同的总训练 FLOPs 下训练。

模型 / 配置	PTB PPL (↓)	WikiText PPL (↓)	AG News PPL (↓)	BoolQ 准确率 (%)	RACE 准确率 (%)	下游 9 项均分 (%)
MDM 基线 (无循环)	108.5 ± 6.61	35.1 ± 1.11	69.3 ± 1.96	46.6 ± 1.1	28.5 ± 0.8	40.0 ± 0.3
LoopMDM (\(S=1\))	124.6 ± 12.02	41.8 ± 1.54	87.8 ± 2.17	53.8 ± 0.8	27.2 ± 0.7	39.4 ± 0.3
LoopMDM (\(S=6\))	100.2 ± 9.44	34.2 ± 1.29	59.2 ± 1.86	53.8 ± 1.8	31.3 ± 0.5	40.7 ± 0.9
LoopMDM (\(S=12\))	90.2 ± 8.59	32.4 ± 1.31	56.6 ± 1.70	57.4 ± 1.3	31.7 ± 0.4	41.4 ± 0.6

在数学推理任务 GSM8K（采用 TinyGSM 11.8M 语料训练的 125M 模型）上，LoopMDM 与同等每步算力的深层非循环基线进行了对决：

模型架构	物理层数	等效每步计算量	GSM8K Top-2 (%)	GSM8K Top-3 (%)
MDM 基线	14 层	基准	74.0	73.8
MDM 深层基线	18 层	1.29× 基准	76.5	76.2
MDM 最深基线	21 层	1.50× 基准 (≈LoopMDM)	79.1	78.5
LoopMDM (\(S=8\))	14 层 (2 层循环)	1.50× 基准	81.8	81.2
LoopMDM (\(S=16\), 外推)	14 层 (2 层循环)	2.07× 基准	82.5	81.9

消融实验¶

在 LM1B 语料上固定训练 FLOPs 预算，针对循环层位置、循环层数 \(n_m\) 以及最大循环轮次 \(S_{\max}\) 展开详细消融（基线 MDM 测试 NLL 为 3.798）：

消融维度	变量设置	测试 NLL (↓)	说明 / 结论
(a) 循环位置 (0-based)	Position 0 (输入端)	3.744	过于靠近输入，浅层特化受损
	Position 1-2 (前中部)	3.729	最优配置，兼顾表征稳定与算子复用
	Position 5 (中间偏后)	3.734	收益略微收窄
	Position 11 (输出端)	3.881	严重劣化，破坏最终分类线性投影
(b) 循环层数 \(n_m\)	\(n_m = 1\)	3.743	循环容量略显不足
	\(n_m = 2\)	3.729	效益/开销比最佳
	\(n_m = 6\)	3.784	接近基线，每步开销偏大
	\(n_m = 12\) (全模型循环)	3.933	严重劣化，训练步数严重被稀释
(c) 最大训练循环轮数 \(S_{\max}\)	\(S_{\max} = 4\)	3.741	循环深度受限
	\(S_{\max} = 12\)	3.729	深度与动态范围最佳折中
	\(S_{\max} = 24\)	3.800	训练时深度过大，随机梯度不稳定

关键发现¶

训练收敛速度倍增：在 LM1B、OpenWebText 和 FineWeb-Edu 三大语料上，LoopMDM 仅需基线 30% 到 43% 的训练 FLOPs（即实现 2.34× 至 3.34× 的计算加速），就能达到标准 MDM 的最终收敛 NLL。
超越同等算力的直接加深（Depth Scaling）：在 GSM8K 数学推理评测中，14 层的 LoopMDM（\(S=8\)）不仅大幅超越同层数的标准 MDM（+7.8 分），更以 81.8% 击败了每步计算量相当的 21 层标准深层 MDM（79.1%），并在外推至 \(S=16\) 时达到 82.5%（+8.5 分提升），证明了在共享权重上反复递归比简单堆叠非共享深层对全局约束任务更有效。
自适应早停实现算力减半无损精简：采用 \(\epsilon=0.1\) 的隐藏层动态早停阈值，模型的平均推理循环次数从固定的 \(S=12\) 骤降至 \(\bar{S} \approx 5.1\)，计算量减少约 50%，但下游任务准确率仅从 41.4% 微调至 41.3%，生成困惑度稳定在 39.11（对比固定 12 步的 38.53）。

亮点与洞察¶

将掩码标记重构为“连续隐式计算草稿纸”：在自回归模型中无法直接引入针对未生成词的并行循环，而 MDM 天然存在的掩码位置为模型提供了固定容量的隐式存储槽位。论文构造的 \(k\)-Clique 理论分离与严格从左到右数独生成实验清晰地证明：循环网络通过在掩码位置间反复交互（Mask-to-Mask Attention），可以在离散输出确定前完成全局一致性修正。
极度克制而优雅的“选择性局部循环”设计：不同于以往盲目追求整体循环导致训练过慢的方案，本文发现仅循环前中部的 2 层即可捕捉全局绝大部分深度收益，不仅维持了输入输出端的高效特化，还极大地控制了训练步数的折损，给后续架构设计提供了极具实操价值的指引。

局限与展望¶

规模限制与超大规模验证空白：当前实验主要基于 125M 至 170M 参数级别的小中型模型，尚未验证该选择性循环策略在数十亿至百亿参数（如 LLaDA 8B/100B）以及超长上下文窗口下的扩展法则与显存开销。
与推理时采样策略的联合探索不足：论文目前主要采用标准去噪采样及 Top-k 解掩码，尚未深入探究循环层与高阶引导策略（如 Remasking、Search-based Decoding、Lookahead Unmasking）在测试时的协同效应。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将循环网络引入掩码扩散语言模型，并从理论与实验双重维度证明了掩码隐状态作为并行工作区的独特性质。
实验充分度: ⭐⭐⭐⭐⭐ 严格遵循 Iso-FLOP 规范，跨三大预训练语料、9 个下游基准、数独推理以及 GSM8K 进行了完整而严密的评测与消融。
写作质量: ⭐⭐⭐⭐⭐ 结构清晰紧凑，逻辑层层递进，从理论分离实例到注意力分析再到自适应推理闭环完整。
价值: ⭐⭐⭐⭐⭐ 为扩散语言模型的架构优化开辟了全新视角，其选择性循环与自适应早停策略兼具极高的实用性与工程借鉴价值。