墨未干时的反复斟酌——LoopMDM如何用"循环"改写AI的训练与推理

论文：Looped Diffusion Language Models
作者：Sanghyun Lee, Chunsan Hong, Seungryong Kim 等
机构：KAIST / KRAFTON / UC Berkeley
arXiv: 2605.26106

一、引子：两个石匠的故事

想象两个石匠雕刻同一座塑像。

第一个石匠有一整条流水线：二十个工匠依次排开，每人负责一道工序——第一个凿轮廓，第二个修比例，第三个打磨细节，依此类推。塑像从头到尾走一遍流水线，就算完成。这是自回归模型（Autoregressive Model, ARM）的工作方式：从左到右，一个字一个字地"刻"，刻完第一个字才能决定第二个字，如同二十个工匠按顺序接力。

第二个石匠的做法完全不同。他先给一块大理石粗略地涂上全身各处的轮廓——像是一个全身都标了记号的半成品，然后让所有工匠同时工作，每个人都在看整体的草图，反复修改自己负责的那一小块。每轮修改后，塑像的全貌就更清晰一点。这是掩码扩散模型（Masked Diffusion Model, MDM）的方式：从一个全是"[待填]"标记的句子开始，逐步去噪、逐步填充。

两个方法都能雕出杰作。但第二个石匠有个奇怪的发现：如果让中间几个工匠反复打磨同一道工序——不是增加人手，而是让同几个人多走几遍——塑像的质量反而比增加整条流水线的新工匠更好。

这就是LoopMDM（Looped Masked Diffusion Model）的核心发现。2025年5月，来自KAIST、KRAFTON和UC Berkeley的研究者证明：在掩码扩散模型中，选择性地循环早期-中间的Transformer层，可以在不增加参数的情况下产生深度扩展的效果，让训练FLOPs减少高达3.3倍，同时在GSM8K数学推理基准上提升8.5个百分点。

更惊人的是，这种"循环"不仅在训练时省成本，还能在推理时按需调整——就像石匠可以根据石料的硬度决定打磨几遍，而不是每次都用固定的工序数。

这个故事，讲的是AI如何从"堆人力"走向"精打细磨"。

🎭 二、从单向书写到全局涂鸦：扩散模型登场

要理解LoopMDM的妙处，得先理解它站在谁的肩膀上。

2.1 自回归模型：多米诺骨牌式的写作

如果你用过ChatGPT，你已经熟悉了自回归模型的体验。你输入一个问题，它从左到右一个字一个字地生成回答，每个新词都依赖于前面所有已经生成的词。

这就像推倒一排多米诺骨牌：第一块倒下会推动第二块，第二块推动第三块。好处是简单、直接、概率链条清晰。坏处也明显——如果前面某一块放歪了，后面的所有骨牌都会跟着歪。在语言模型中，这叫错误累积（error accumulation）：开头一个词选错了，后面的整个句子都可能被带偏。

更深层的问题是，自回归模型在生成每个词时，只能"看"到已经生成的部分。它像一个只能从左往右写的作家，永远不知道句子的后半部分会是什么，只能根据前文猜测。这种"因果掩码"（causal masking）既是它的力量来源（保证生成过程的连贯性），也是它的天花板——因为人类写作从来不是纯粹的从左到右过程。你会先写一个大致框架，再回头修改开头；你会在写完结尾后才发现中间需要插入一个过渡句。

2.2 扩散模型：从噪声中浮现的图像

扩散模型的灵感来自一个反直觉的想法：如果你有一张照片，你可以一步一步往里面加噪声，直到它变成完全的随机雪花屏；然后，如果你训练一个神经网络学会这个"加噪"过程的逆过程，它就能从纯噪声中一步一步恢复出原始图像。

这就像一个魔术师把一幅画撕碎、烧毁、扔进搅拌机，然后反向操作，让画从灰烬中重新拼合出来。2020年，DDPM（Denoising Diffusion Probabilistic Models）让这种方法在图像生成领域爆发，催生了DALL-E、Stable Diffusion等现象级应用。

扩散模型的魅力在于它的迭代精炼（iterative refinement）。它不是一次猜对答案，而是通过多轮去噪逐步逼近真实。每一轮都基于当前的状态做局部改进，像雕塑家从粗糙的泥胚开始，每轮都比上一轮更像最终作品。

2.3 掩码扩散模型：把"扩散"带进语言世界

但图像是连续的，语言是离散的。你不能把一句话"逐渐变成噪声"——词语要么存在，要么被删掉（掩码），没有中间态。

2021年，Austin等人在《Structured Denoising Diffusion Models in Discrete State-Spaces》中提出了一种巧妙的方案：掩码扩散。模型从一个完全被[MASK]标记遮蔽的句子开始。在每一步"去噪"中，模型观察当前的句子（部分是真实词，部分是掩码），然后预测每个掩码位置应该填什么词。预测完成后，部分掩码被揭开，露出新生成的词，但另一些位置仍然保持掩码，留到后面的步骤处理。

这就像你在做填空题，试卷上所有空都是空着的。第一轮，你凭直觉填了一些最有把握的空；第二轮，你用已经填好的信息去推断剩下的空；第三轮，你回头检查前面填的是否正确，因为新的信息可能让你改变之前的判断。

这就是MDM的核心：掩码位置不仅是"待填的空白"，它们还充当了并行计算的工作区（parallel workspace）。一个被掩码的位置可以在多轮迭代中存储、修改、交换信息，这是自回归模型中不存在的自由度。ARM的每个位置一旦生成就被"锁定"，永无反悔之日；MDM的掩码位置则像黑板上的草稿，可以擦了重写。

正是这个特性，让循环层在MDM中产生了独特的化学反应。

➿ 三、循环层：为什么"反复"比"加深"更聪明

3.1 循环Transformer的前世

"循环"（looping）在Transformer架构中不是新概念。2018年，Dehghani等人提出Universal Transformer，让一个共享的Transformer块反复应用于输入，模拟图灵机的迭代计算。2023年，Giannou等人证明，循环Transformer可以被编程为通用计算机，理论上能执行任意算法。

在自回归模型中，循环已经被探索为一种测试时计算扩展（test-time compute scaling）的手段：通过增加循环次数，让模型在推理时"想得更深"。但此前的研究主要关注参数效率（用更少参数达到同等效果）或推理时的精炼，对于训练效率的提升一直不理想——说白了，循环在自回归模型里是个"锦上添花"的装饰品，不是"雪中送炭"的效率利器。

3.2 为什么MDM是循环的更佳土壤

MDM天然具有两层迭代结构：外层是跨时间步的去噪迭代（denoising steps），内层是每个时间步中的前向传播（forward pass through layers）。研究者最初怀疑：既然MDM已经在时间维度上迭代了，再加一个层级的循环会不会是冗余的？

答案是：不会冗余，反而协同。

关键原因在于MDM的"掩码工作区"。在ARM中，循环Transformer反复处理的是同一组已生成的token，信息 content 基本固定，循环的收益主要来自"深度思考"。但在MDM中，循环层处理的是包含大量掩码位置的中间表示。每次循环，掩码位置之间可以交换信息、修正彼此的预测、建立全局一致的结构。这种"掩码位置之间的交互"（interactions among masked positions）是ARM根本无法实现的。

想象一个解谜游戏：ARM是一个玩家按顺序打开房间的门，开完一扇才能开下一扇，开了就不能回头；MDM是一个玩家同时面对所有上锁的门，每轮可以尝试打开所有门，没打开的还是锁着，但已经从打开的门里获得了线索。循环层则让这位玩家在同一批门前来回审视多遍——第一遍打开容易的门，第二遍用新信息打开难一点的门，第三遍检查之前开的门是否互相矛盾。这种"反复审视"对于需要全局一致性的任务（如数学推理）至关重要。

3.3 "早期-中间层"的选址智慧

LoopMDM最精妙的设计决策之一是在哪里循环。论文做了系统性的消融实验，测试了在不同位置循环的效果：

这个结果揭示了一个深刻洞见：循环最有效的时机，是在"原始信息已经初步编码"但"还没有进入最终输出特化"的那个窗口。太早，循环的是毫无语义的原始嵌入；太晚，循环的是已经为预测任务过度调优的表示，没有空间再调整。中间层恰好处于"半成品"状态——既有足够的结构化信息可供操作，又保留了灵活性。

这就像厨师做菜：刚把食材切好就反复炒没有意义（还没调味），快出锅了再反复炒也不合适（已经定型），而在食材下锅、初步受热但尚未定型的那个阶段多加翻炒，最能提升口感。

3.4 循环次数S的魔法

LoopMDM在训练中采用随机循环策略：每步训练时，从均匀分布U{1, ..., S_max}中随机采样循环次数S。这使得共享的循环块必须适应各种"有效深度"——浅时（S=1）要像浅层网络一样工作，深时（S=12）要像深层网络一样工作。这种训练方式强迫循环块学习一个深度自适应的算子（depth-adaptive operator）。

实验揭示了一个有趣的现象：

• S=1时，LoopMDM反而差于基线。这是因为共享权重在单次前向传播中无法展现其迭代潜力，表现得像一个被强行压缩的浅层模型。
• S=6时，捕获了大部分收益。适度循环就足以激活共享块的深层能力。
• S=12（训练最大值）时，达到最佳性能。
• S=24（超出训练范围）时，性能继续改善。这意味着学到的算子可以泛化到训练时从未见过的深度——这是循环架构独有的特性。固定深度的非循环模型永远无法在推理时"自动变深"。

📊 四、数字说话：3.3倍效率与8.5点提升的真相

4.1 预训练效率：用更少算力追上同等性能

论文在三个标准语言建模语料库上评估了LoopMDM的训练效率：

在所有三个数据集上，LoopMDM在S≥6时一致优于基线。LM1B上的3.34倍是最惊人的——LoopMDM用不到三分之一的训练算力就追平了同等规模MDM的最终性能。

这里的关键是匹配训练总FLOPs的实验设计。基线MDM训练1M步，batch size 512。LoopMDM的循环增加了每步的FLOPs，因此研究者按比例减少了训练步数，确保总训练FLOPs与基线相同。在这种公平的比较下，LoopMDM仍然胜出。换句话说，同样的电费，循环模型学到的更多。

4.2 数学推理：GSM8K上的8.5点飞跃

如果语言建模的改进只是"省点电费"，那GSM8K的结果就是"亮出肌肉"。

GSM8K是一个小学数学应用题数据集，测试模型的多步推理能力。LoopMDM在这个任务上的表现堪称惊艳：

• LoopMDM S=16（超过训练最大值S_max=8）在Top-2解码策略下，比同参数14层MDM基线提升+8.5个百分点
• 在Top-3解码策略下提升+8.1个百分点
• LoopMDM仅用基线43%的训练FLOPs就追平了基线的最终准确率

更关键的是对比实验：研究者还训练了更深的非循环MDM（18层和21层），其每步训练FLOPs与LoopMDM相当。结果？

• LoopMDM S=8 超过21层非循环MDM +2.7点
• LoopMDM S=16 超过21层非循环MDM +3.4点

这个结论掷地有声："重复应用一个共享块"比"简单地堆叠更多层"对推理任务更有效。21层模型的参数更多，训练成本更高，但在推理能力上却输给了循环两次中间层的LoopMDM。

为什么数学推理特别受益于循环？论文的注意力分析提供了线索。

4.3 注意力分析：掩码位置之间的秘密对话

研究者测量了循环层中掩码到掩码的注意力（mask-to-mask attention）——即查询来自掩码位置、键也来自掩码位置的平均注意力质量。

结果清晰得像一个故事：

• S=1时，mask-to-mask注意力很低——掩码位置之间几乎没有信息交换。
• S=6时，注意力中等——循环开始激活掩码位置的交互。
• S=12时，注意力接近饱和——掩码位置之间建立了密集的信息通道。
• S=24（超出训练范围）时，维持在高水平——学到的交互模式泛化到新深度。

这直接支持了"掩码作为工作区"的理论：循环层让掩码位置从被动的"等待填充"变成了主动的"信息交换节点"。在数学推理中，这相当于草稿纸上的多个未知数可以互相参考、互相约束——比如"如果x=3，那y必须等于7才能满足方程；但如果y=7，那么z就不能是5"。这种交叉约束在自回归模型中只能通过生成顺序间接实现，而在循环MDM中，它在每一轮去噪步骤内部就能发生。

研究者还发现，循环的收益在去噪过程中不均匀分布：中间时间步收益最大（既有足够的上下文信息，又有足够多的未解析位置需要交互），早期（信息不足）和晚期（几乎已解决）收益较小。这一发现启发了后续的自适应循环策略。

🧩 五、Sudoku实验：当推理被剥夺选择权时

为了进一步隔离循环对"单步内部计算"的作用，研究者设计了一个精妙的Sudoku实验。

5.1 实验设计的巧思

标准MDM解Sudoku时可以使用自适应解掩码顺序：先解最容易的格子，用已确定的值去约束其他格子。这种策略甚至可以被7B规模的自回归模型使用（虽然对它们来说仍然很难）。

但研究者把生成顺序强制固定为从左到右，不管每个格子的难度。在这种苛刻的条件下，任何性能提升都必须来自单步去噪过程中的内部计算，而不是来自聪明的生成顺序。

这就像让一个棋手蒙眼下棋，不让他选择先看哪里，必须按固定顺序审视棋盘——如果他还能赢，那赢的一定是纯粹的计算力。

5.2 结果：三遍循环，从24个错误到0

实验在100k Sudoku数据集上训练了一个极小的模型（1层，384维，12头）及其循环版本（循环同一层，S_max=6）。

结果令人震撼：

可视化错误轨迹，研究者看到了循环的"思考过程"：

• 第一轮（S=1）：模型给出局部合理但全局不一致的赋值——每个格子看起来都对，但放在一起就矛盾（比如同一行有两个相同的数字）。
• 第二轮（S=2）：模型开始修正一些冲突，但仍有7处未解决。
• 第三轮（S=3）：掩码位置之间的交互最终消除了所有全局不一致，得到完美解答。

这个实验的最重要结论是：循环让掩码位置在token被最终确定之前有机会修订全局不一致的早期预测。在自回归模型中，一旦一个词被生成，就永远锁死；在循环MDM中，掩码位置可以在多轮循环中互相"谈判"，直到达成一致。

⚡ 六、推理时自适应：按需分配算力

LoopMDM的第二个关键优势是推理时的灵活计算扩展。这得益于训练时的随机循环策略——模型学会了在各种循环深度下都工作良好，因此推理时可以根据需要调整S。

6.1 自适应循环策略

固定循环次数并不总是最优的。论文提出了一种基于隐藏状态稳定性的自适应停止准则：

对于每个去噪时间步 t:
    对于每次循环 k = 1, 2, ..., S_max:
        计算 H^(k) = mid-block(H^(k-1))
        计算相对变化: ||H^(k) - H^(k-1)|| / ||H^(k)||
        若 相对变化 < ε (阈值) 或 k = S_max:
            停止循环，输出到Tail

当隐藏状态在连续两次循环之间几乎没有变化时，说明模型已经"想透"了，不需要再循环。

6.2 自适应策略的效果

自适应策略将平均循环次数从12降至5.1-5.2，计算量减少了约57%，而下游准确率几乎没有下降（41.4% → 41.3%）。

更有趣的是，自适应策略自然地将更多循环分配给中间时间步，早期和晚期时间步分配较少——这与注意力分析中发现的"中间时间步收益最大"完全吻合。这说明模型自己学会了在哪里"多想一想"，哪里"差不多就行"。

🔮 七、对AI效率革命的启示

LoopMDM的意义远不止一篇技术论文。它指向了一个更深刻的范式转变。

7.1 "堆参数"时代的终结

过去五年，大型语言模型的进步主要遵循一条简单粗暴的路径：更大参数 + 更多数据 + 更多GPU。GPT-3有1750亿参数，GPT-4据传超过万亿，每个新版本都在参数数量上碾压前一个。这场军备竞赛烧掉了数十亿美元的算力，也让小公司和个人研究者越来越难以参与。

LoopMDM展示了一条截然不同的路径：在固定参数预算内，通过更聪明的架构设计实现深度扩展。3.3倍的训练效率提升意味着同样的模型可以用三分之一的GPU时间训练到同等水平，或者用同样的资源训练出显著更好的模型。在GSM8K上，一个125M参数的LoopMDM打败了参数更多、训练更深的非循环模型。

这不是渐进式改良，这是架构层面的效率跃迁。

7.2 "推理时计算"成为一等公民

传统上，模型的训练和推理被割裂看待：训练时拼命堆资源得到一个固定模型，推理时只能被动接受模型的输出质量。但LoopMDM让推理过程本身变成了可扩展的计算资源——增加循环次数就能提升推理质量，而且自适应策略让这种扩展是按需的。

这与AI领域近期的另一大趋势——测试时计算扩展（test-time compute scaling）——形成呼应。DeepSeek-R1、OpenAI o1/o3系列都在展示：让模型在推理时"多想一想"，可以在不增加训练成本的情况下大幅提升复杂任务的表现。LoopMDM为这一趋势提供了新的技术路径：不是通过更长的思维链或更多的采样，而是通过架构层面支持的迭代精炼。

7.3 循环作为"深度"的新定义

在深度学习的历史上，"深度"几乎等同于"层数"。ResNet有152层，Transformer-base有12层，更大的模型有96层甚至更多。但LoopMDM挑战了这个定义：真正的深度不一定是物理层的堆叠，而是有效计算的迭代次数。

一个12层模型循环12次，其有效深度是144层"计算等价"，但参数只有12层的量。这种"虚拟深度"带来了参数效率的革命——它让小型模型可以拥有大型模型的计算能力，而不需要大型模型的内存占用。

对于边缘设备、手机端AI、实时应用，这意味着原本不可能的运行大型模型推理，现在或许可以通过循环架构在有限硬件上实现。

7.4 掩码工作区的启示

LoopMDM的成功根植于MDM的一个独特属性：掩码位置作为并行计算工作区。这提示了一个更广泛的原理——显式的"未决状态"可能是推理能力的关键。

自回归模型没有真正的"未决状态"：每个位置一旦生成就被锁定。人类思维则完全不同：我们在思考时会在工作记忆中保留大量"待确认"的信息，让它们互相碰撞、修正、重组，直到一个一致的结构浮现。MDM的掩码位置正是这种"认知工作区"的计算类比。

LoopMDM通过循环放大了这个工作区的效果。也许未来的AI架构设计应该更主动地引入类似"草稿空间""工作记忆""暂存缓冲区"的结构，而不是一味追求端到端的确定性流水线。

八、结语：打磨的哲学

回到开头两个石匠的故事。

LoopMDM教给我们的不是某个具体的技术技巧，而是一种关于"如何做深度工作"的哲学：与其增加更多工匠，不如让已有的工匠多打磨几遍。深度不是人数的函数，而是专注与反复的函数。

在AI领域，这意味着我们或许正在接近一个拐点——从"大力出奇迹"的蛮力时代，进入一个"精打细磨"的效率时代。3.3倍的训练效率提升不是小数字，它意味着同样的碳足迹可以训练三个同等模型，或者一个模型用三分之一的能源达到同等水平。在一个算力即权力、电力即约束的世界里，这种效率革命可能比参数数量的军备竞赛更有意义。

对于步子哥这样的AI内容创作者，LoopMDM的启示也很直接：下一个爆款模型可能不是来自更疯狂的数据中心，而是来自某个KAIST实验室里几个研究者对"循环"这一古老想法的重新发现。有时候，最深刻的创新不是发明新东西，而是在新土壤上让旧种子重新发芽。

掩码扩散是土壤，循环是种子，而3.3倍的效率提升和8.5点的GSM8K飞跃，就是开出的花。

参考文献

[1] Anil, C., et al. (2022). Path independent equilibrium models can better exploit test-time computation. NeurIPS.

[2] Austin, J., et al. (2021). Structured denoising diffusion models in discrete state-spaces. NeurIPS.

[7] Chao, C.-K., et al. (2025). Beyond masked and unmasked: Discrete diffusion models via partial masking. arXiv:2505.18495.

[11] Cobbe, K., et al. (2021). Training verifiers to solve math word problems. arXiv:2110.14168.

[15] Dehghani, M., et al. (2018). Universal transformers. arXiv:1807.03819.

[17] Deschenaux, J., et al. (2025). Partition generative modeling: Masked modeling without masks. arXiv:2505.18883.

[18] Fan, A., et al. (2024). Looped transformers for length generalization. arXiv:2409.15647.

[19] Geiping, J., et al. (2025). Scaling up test-time compute with latent reasoning: A recurrent depth approach. arXiv:2502.05171.

[20] Giannou, A., et al. (2023). Looped transformers as programmable computers. ICML.

[27] Hong, C., Lee, S., & Ye, J. (2026). Unifying masked diffusion models with various generation orders and beyond. arXiv:2602.02112.

[37] Liu, B., et al. (2023). TinyGSM: achieving >80% on GSM8K with small language models. arXiv:2312.09241.

[38] Lou, A., Meng, C., & Ermon, S. (2024). Discrete diffusion modeling by estimating the ratios of the data distribution.

[46] Nie, S., et al. (2025). Scaling up masked diffusion models on text. arXiv:2410.18514.

[47] Nie, S., et al. (2025). Large language diffusion models.

[50] Peebles, W., & Xie, S. (2023). Scalable diffusion models with transformers.

[Lee et al., 2026] Lee, S., Hong, C., Kim, S., Lee, J., Park, J., & Park, D. (2026). Looped Diffusion Language Models. arXiv:2605.26106.

#LoopMDM #掩码扩散模型 #Transformer #AI效率 #推理时计算扩展 #论文解读 #费曼风格
#论文 #arXiv #AI #扩散模型 #小凯 #每日论文