🤖 machine learning

Recursive Scaling in Masked Diffusion Models

本文介绍了递归掩码扩散模型（Recursive Masked Diffusion Models, R-MDMs），该模型通过增加递归作为第三个缩放维度，允许单个 Transformer 迭代地细化输出，从而在提升参数效率和推理速度的同时，以更少的参数和去噪步骤实现了与规模大得多的非递归模型相当的性能。

原作者： Alba Carballo-Castro, Julianna Piskorz, Paulius Rauba, Mihaela van der Schaar, Pascal Frossard

发布于 2026-06-17

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Alba Carballo-Castro, Julianna Piskorz, Paulius Rauba, Mihaela van der Schaar, Pascal Frossard

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

核心思想：更聪明的循环，而非更大的大脑

想象一下，你正在尝试解决一个困难的谜题，比如数独或数学题。通常，当计算机遇到瓶颈时，标准的建议是：“构建一个更大的大脑。” 在 AI 领域，这意味着增加参数（让模型变得更大、更复杂）或者让它思考更长的时间（更多的步骤）。

这篇论文提出了一种不同的方法。他们不是在构建一个更大的大脑，而是教同一个大脑连续对问题进行多次思考，通过每一次迭代来完善它的答案。他们称之为递归缩放（Recursive Scaling）。

可以这样理解：

传统方式（向上缩放）： 你雇佣了 30 位不同的专家来看一次这个谜题。
递归方式（Recursion）： 你雇佣了一位专家，但你让他看一眼谜题，做出一个猜测，退后一步，观察自己的猜测，纠正它，再看一遍，再次纠正。他这样重复 5 次或 10 次。

研究发现，通过循环工作的单个专家，其表现往往优于那 30 人的专家团队，尽管这位单人专家的规模要小得多，运行成本也更低。

工作原理：“草拟与润色”类比

为了理解这项技术，让我们看看这些 AI 模型（被称为掩码扩散模型，Masked Diffusion Models）通常是如何工作的：

掩码游戏： 想象一个句子，其中一些单词被隐藏了（被掩盖了）。AI 的任务是同时猜出缺失的单词。
标准流程： AI 做出一个猜测。然后，它揭开（un-mask）几个单词，再次进行猜测。它重复这个过程很多次（就像剥洋葱一样，一层一层地剥开），直到整个句子显现出来。

创新之处：
作者意识到，在这些层级中的每一个层级内，AI 都可以进行一次“快速草拟与润色”环节。

标准 AI： 做一个猜测 $\rightarrow$ 剥开一层 $\rightarrow$ 做一个新的猜测。
递归 AI (R-MDM)： 做一个猜测 $\rightarrow$ 观察这个猜测，批判它，并将其改进 5 次 $\rightarrow$ 然后再剥开一层。

他们称之为递归深度（Recursive Depth）。这就像是在模型提交答案之前，给它一次“第二次思考”或“第三次思考”的机会。

为什么这是一个游戏规则的改变者

论文在三种类型的任务上进行了测试：数独（逻辑谜题）、Countdown（数学谜题）和 Text8（文本写作）。以下是他们的发现：

1. “小脑”在逻辑任务中胜出

在像数独和 Countdown 这样的结构化谜题上，递归模型表现得非常出色。

结果： 一个循环 5 次的小模型，其表现竟然与一个规模大 5 倍但只循环一次的庞大模型不相上下（甚至更好）。
类比： 这就像是一个能在脑海中预判 10 步棋的棋手，对比一个虽然是大师但只能思考 1 步棋的棋手。这种“循环”让小模型能够看到全局并修正错误，而大模型仅仅是做出了一个自信但可能错误的猜测。

2. 加速过程

通常，为了得到完美的答案，你需要运行 AI 经过许多“去噪步骤”（即许多层猜测）。

结果： 递归模型在更少的步骤内就达到了高质量的答案。
类比： 想象你在编辑一份文档。
- 普通 AI： 写一段话，停下来，编辑一下，停下来，再编辑一下。它需要 40 轮编辑才能达到完美。
- 递归 AI： 写一段话，然后立即连续阅读并编辑 3 次，然后再 进入下一段。它只需 15 轮就能达到同样的完美质量。
- 益处： 这节省了大量的计算能力（时间与电力）。

3. 局限性：取决于任务

论文指出，这种技巧最适用于结构化问题（如数学和逻辑谜题），因为这些问题有明确的规则和依赖关系。

Text8 的结果： 当他们尝试用这种方法处理随机文本写作（如维基百科文章）时，递归模型在标准数学得分上实际上比大模型表现得更差。
结论： “循环”这一超能力在解决需要根据规则检查工作的谜题时非常有效。对于创意写作，单纯把模型做大可能仍然是更好的选择。

“魔法”总结

这篇论文介绍了一种新的 AI 缩放方式，它不仅仅意味着“把它变大”。

旧规则： 要想变得更聪明，增加更多层（更多参数）。
新规则： 要想变得更聪明，增加更多循环（让模型完善自己的工作）。

通过让一个较小的模型通过重新处理自己的输出来进行“深度思考”，作者实现了：

在逻辑谜题上获得更好的性能。
用更少的步骤获得高质量答案。
更低的成本，因为他们不需要构建庞大且昂贵的模型。

简而言之：不要只是构建一个更大的大脑；要教会大脑如何反复检查自己的工作。

技术摘要：掩码扩散模型中的递归缩放

问题陈述

掩码扩散模型（Masked Diffusion Models, MDMs）已成为一种极具前景的离散序列生成范式，其提供的全并行生成能力与自回归（AR）模型的顺序生成特性形成了鲜明对比。然而，将 MDMs 扩展以实现高质量生成的传统路径依赖于两个维度：增加参数量（宽度/深度）或增加去噪步数（ $T$ ）。这种方法在训练（更大的模型）和推理（更多的前向传播）方面都会产生显著的计算成本。虽然递归计算（多次应用相同的变换）在自回归模型和监督学习中已被证明在提高参数效率方面具有潜力，但在 MDMs 的背景下仍未得到充分探索。本文解决的核心挑战在于：递归深度是否可以作为第三个正交的缩向维度，在不增加参数量的情况下提升生成质量，并且能否替代额外的去噪步数以减少推理时的计算量。

方法论

作者引入了递归掩码扩散模型（Recursive Masked Diffusion Models, R-MDMs），该模型在标准的 MDM 训练范式中引入了递归循环机制。

架构

R-MDM 不再仅单次应用 $K$ 层 Transformer 层，而是使用一个共享的 $K$ 层 Transformer 块 $f_\theta$ ，在单个去噪步骤内连续应用 $L$ 次。

参数效率： 尽管具有 $K \times L$ 的有效深度，但 R-MDM 的参数量与单次通过的 $K$ 层模型相同。
前向传播： 给定噪声输入 $x_t$ ，隐藏状态被迭代更新： $h^{(\ell)} = f_\theta(h^{(\ell-1)}, \ell, L)$ 。第 $\ell$ 次循环的输出作为第 $\ell+1$ 次循环的输入。
步进嵌入（Step Embedding）： 为了让模型能够区分不同的递归阶段，模型引入了步进嵌入 $v^{(\ell)}$ 。该嵌入通过归一化后的循环进度 $s_\ell = (\ell-1)/(L-1)$ 经由一个小型的 MLP 推导得出，使模型能够根据递归深度进行专门化计算。

训练目标

论文评估了四种针对递归循环的监督策略：

全步监督（ALL）： 在所有 $L$ 个循环中平均交叉熵损失，提供密集的梯度信号（受深度监督学习启发）。
最终步监督（FINAL）： 仅监督最后一次循环的 Logits，这是对标准 MDM 目标的扩展。
加权监督（Weighted）： 为不同循环的损失分配非均匀权重。
截断监督（Truncated）： 仅监督最后 $k$ 个循环。

作者发现，全步监督（All-steps loss） 的表现最为稳健，尤其是在低计算量的情况下。

推理与调度

在推理时，模型可以使用与训练深度 $L_t$ 不同的采样递归深度 $L_s$ 进行评估。作者探索了多种 $L_t$ 训练方案，包括固定深度、随机采样和基于课程的学习调度。研究发现，反向课程（Reverse Curriculum）（从高递归深度开始并逐渐降低到 1）表现更为优异，它迫使模型在训练早期就学习高容量的迭代细化策略。

核心贡献

引入递归深度作为缩放维度： 论文提出在去噪步骤内进行递归循环，作为在不增加参数量的情况下增加有效模型深度的机制。
展示参数效率： 研究表明，R-MDMs 在结构化推理任务上能匹配甚至超越参数量显著更大的非递归基准模型（高达 5 倍参数量）。
推理效率： 研究证明，递归细化可以部分替代去噪步数，使得模型能够以更少的总前向传播次数（ $T \times L$ ）达到高质量生成。
任务相关性分析： 本工作系统地评估了在结构化任务（数独、Countdown）和非结构化任务（Text8）中，递归、参数量与去噪步数之间的权衡关系。

实验结果

实验在三个数据集上进行：Sudoku（9x9 和 25x25）、Countdown（包含 3、4 和 5 个操作数的算术规划）以及 Text8（字符级语言建模）。

结构化推理任务（Sudoku & Countdown）

参数替代（RQ1）： 递归模型 $(K \otimes L)$ $(K \otimes L)$ 一致地匹配或优于等 FLOPs 的单次通过基准模型 $(KL \otimes 1)$ $(K L \otimes 1)$ 。
- 示例： 在 Sudoku 9x9 上，一个 $(6 \otimes 5)$ 模型（10.6M 参数）在 $T=5$ 时达到了 93.6% 的有效谜题率（VPR），这匹配了一个 30 层单次通过模型（53.1M 参数）的表现，而后者在相同步数下仅为 77.8%。
- 示例： 在 Countdown-4 上，一个 $(3 \otimes 5)$ 模型（5.5M 参数）在 $T=20$ 时达到了 70.2% 的目标达成率（RTR），超过了 15 层基准模型（26.7M 参数）的 69.4%。
步数减少（RQ2）： 递归模型需要更少的去噪步数来达到目标质量阈值。
- 一个循环模型（ $L_s=3, T=5$ ）通过 15 次总前向传播即可达到 95% 的有效谜题率，而同等的单次通过模型则需要 4 步（即 40 次前向传播），实现了 2.7 倍的计算量削减。
- 在 80% 掩码率的 Sudoku 25x25 任务中，递归模型（ $L=5$ ）在 $T=1$ 时达到了 91.8% 的 VPR，而基准模型在此步数下几乎完全失败。
灵活性： 使用反向课程（ $L_t$ 递减）训练的模型可以有效地使用不同的采样深度 $L_s$ 进行评估，从而在推理时实现灵活的计算分配。

非结构化语言建模（Text8）

在 Text8 数据集上，递归模型在基于似然性的指标（NLL、Perplexity）方面表现不如单次通过基准模型。
然而，定性分析表明，尽管似然得分较低，递归模型生成的文本更具连贯性，这表明标准指标在此类场景下可能具有误导性。作者指出，递归的优势在具有明确迭代结构的领域中最为显著。

意义与主张

论文声称，递归深度是 MDMs 中一种具有实际用途的缩放机制，特别是在结构化生成任务中。其意义主要体现在两个方面：

参数效率： R-MDMs 能以极小的参数量实现大幅领先于其规模的非递归模型性能，使得在资源受限的环境下进行高质量生成成为可能。
计算效率： 通过用递归细化替代额外的去噪步数，R-MDMs 可以减少推理时所需的总前向传播次数，优化生成质量与延迟之间的权衡。

作者总结道，虽然递归是结构化领域中执行迭代约束传播的强大机制，但其在非结构化语言建模中的效用仍是一个需要进一步研究的开放性问题。这项工作丰富了人们对生成模型缩放（超越传统的以参数量为中心的视角）的理解。