💬 NLP

AdaPonderLM: Gated Pondering Language Models with Token-Wise Adaptive Depth

本文提出了 AdaPonderLM，一种基于门控机制的自监督循环语言模型，它通过在预训练阶段学习每个 Token 的自适应退出策略并引入 KV 状态复用机制，在保持语言建模困惑度与下游任务性能不变的同时，显著降低了推理计算成本。

原作者： Shixiang Song, He Li, Zitong Wang, Boyi Zeng, Feichen Song, Yixuan Wang, Zhiqin John Xu, Ziwei He, Zhouhan Lin

发布于 2026-03-12

📖 1 分钟阅读☕ 轻松阅读

原作者： Shixiang Song, He Li, Zitong Wang, Boyi Zeng, Feichen Song, Yixuan Wang, Zhiqin John Xu, Ziwei He, Zhouhan Lin

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

这是一篇关于让 AI 变得更“聪明”且更“省力”的论文。为了让你轻松理解，我们可以把大语言模型（LLM）想象成一个正在写文章的超级作家，而这篇论文提出的 AdaPonderLM 就是给这位作家配备的一套**“智能思考节奏控制器”**。

🧠 核心问题：以前的作家太“死板”了

想象一下，你让一个作家写一篇文章。

以前的做法（固定深度）： 无论写的是“今天天气真好”这种简单句子，还是“量子力学如何影响宏观宇宙”这种难句子，作家都被强制要求必须反复思考 10 遍才能写下每一个字。
- 后果： 写简单句子时，他在第 3 遍就想通了，但为了凑够 10 遍，他不得不假装继续思考，浪费了大量脑细胞（计算资源）；写难句子时，10 遍可能还不够，但他被迫停笔了。
- 比喻： 就像让一个小学生和一位数学家解同一道数学题，规定两人都必须算满 100 步。小学生算 3 步就懂了，后面 97 步纯属浪费时间；数学家算 100 步可能还没算完。

💡 解决方案：AdaPonderLM（自适应思考者）

这篇论文提出的 AdaPonderLM，给作家装了一个**“智能红绿灯”**。

1. 动态调整思考时间（Token-wise Adaptive Depth）

简单句子（容易的词）： 作家思考了一下，发现“今天天气真好”很简单，“智能红绿灯”立刻变绿，告诉他：“停！你已经想明白了，直接写下来，不用继续纠结了。”
- 效果： 省下了后面 9 次思考的时间。
难句子（复杂的词）： 作家思考了几遍，发现“量子力学”很难懂，“红绿灯”一直亮着红灯，告诉他：“继续想！还没想透，再深入一层。”
- 效果： 把省下来的时间，全部用在那些真正需要深度思考的难点上。

2. 怎么保证“想停就停”？（门控机制与 KV 复用）

这里有两个关键技术，我们可以用**“接力赛”和“存折”**来比喻：

门控机制（Gated MLP）： 这是一个**“裁判”**。每思考一轮，裁判就检查一次：“这个词还需要再想吗？”如果不需要，就发出“停止”指令。
- 创新点： 以前的裁判是死板的（不管谁，都喊停或不停），而 AdaPonderLM 的裁判是**“看人下菜碟”**的，每个词都有专属的裁判，而且是在训练过程中自己学会怎么当裁判的，不需要人类手把手教。
KV 复用机制（KV Reuse）： 这是最巧妙的地方。
- 当作家决定“这个词不用想了”时，他不需要把之前思考的所有笔记（Key/Value 状态）扔掉，而是直接盖上“已归档”的印章，原封不动地保留。
- 比喻： 就像你在写代码，遇到一个已经算好的常量，你不需要重新运行一遍函数，直接**调用缓存（Cache）**里的结果就行。这样既保证了逻辑连贯，又极大地节省了算力。

📊 实验结果：省了钱，没降质

研究人员在从很小（70M 参数）到很大（2.8B 参数）的模型上都测试了这套方法：

省算力： 在保持文章质量（困惑度 PPL）几乎不变的情况下，推理时的计算量减少了约 10%。
- 通俗说： 以前写 100 个字要跑 100 次引擎，现在平均只需要跑 90 次，但写出来的文章一样好。
更聪明： 模型真的学会了“把力气用在刀刃上”。分析显示，那些让模型感到“困惑”（高 NLL）的词，模型确实分配了更多的思考轮次；而简单的词，很快就通过了。
下游任务强： 在回答常识问题、阅读理解等任务上，它的表现和那些“死板”但计算量大的模型一样好，甚至更好。

🚀 总结：这意味着什么？

AdaPonderLM 的核心思想就是：不要平均用力，要因材施教。

以前： 无论问题难易，AI 都“一刀切”地思考固定次数。
现在： AI 学会了自我评估。简单的词“秒懂”直接过，难的词“深思熟虑”多花点时间。

这就好比一个高效的团队：

遇到简单任务，大家快速完成，不磨蹭；
遇到棘手难题，大家集中火力，反复推敲。

最终结果是：团队（AI）干活更快了，成本更低了，但产出的质量（文章/答案）依然顶尖。 这对于未来让 AI 在手机上、在普通电脑上跑得更快、更省电，有着巨大的意义。

以下是关于论文 AdaPonderLM: Gated Pondering Language Models with Token-Wise Adaptive Depth 的详细技术总结：

1. 研究背景与问题 (Problem)

测试时扩展（Test-time Scaling）的局限性：近年来，通过循环（Recurrent）或迭代（Iterative）Transformer 架构在推理阶段分配更多计算资源以提升大语言模型（LLM）性能成为热点（如 Universal Transformer, PonderLM 等）。然而，现有的大多数预训练循环模型在实际应用中通常运行固定次数的迭代。
资源分配不均：固定迭代次数导致计算资源的低效分配：
- 简单 Token：被分配了过多的计算（冗余迭代），造成浪费。
- 复杂 Token：可能因迭代次数不足而得不到充分处理。
现有方法的不足：
- 早期的自适应计算时间（ACT）和早退（Early Exit）机制多依赖人工设计的阈值或监督信号。
- 现有的自适应方法往往需要在预训练后引入，依赖硬标签，或仅支持静态剪枝，缺乏在完全自监督预训练阶段实现的端到端 Token 级自适应深度机制。

2. 核心方法论 (Methodology)

作者提出了 AdaPonderLM，一种基于自监督学习的循环语言模型，能够在预训练过程中自动学习每个 Token 的“思考”深度（即迭代次数），无需人工调整剪枝比例。

2.1 模型架构与门控机制

基础框架：基于 PonderLM 架构，使用共享的 Transformer 骨干网络进行循环 refinement。
迭代特定门控（Iteration-specific MLP Gates）：
- 在每次迭代 $i$ 中，使用一个独立的 MLP 层根据当前隐藏状态 $h^i$ 生成每个 Token 的门控 Logits。
- 通过 Sigmoid 函数将 Logits 映射为停止概率 $s^i$ 。
- 设计动机：实验发现共享单一 MLP 会导致训练不稳定，因此为每次迭代使用独立的 MLP。
单调剪枝掩码（Monotonic Halting Mask）：
- 引入一个持久化的二元掩码 $m^i$ 。当门控概率低于阈值 $\tau$ （如 $10^{-4}$ ）时，Token 被标记为“剪枝”（停止迭代）。
- 一旦 Token 被剪枝，它在后续迭代中保持剪枝状态，确保深度单调递减。
更新机制：利用 PonderLM 的概率加权更新策略，将词汇表嵌入的期望值作为候选更新，并通过门控残差连接更新 Token 表示。

2.2 训练 - 推理一致性机制 (KV Reuse)

这是实现实际加速的关键创新：

问题：在训练时，如果 Token 被剪枝，其 KV Cache 在后续迭代中不应重新计算，否则会导致训练与推理行为不一致。
解决方案：提出 KV 复用机制。
- 对于被剪枝的 Token，在后续迭代中直接复用上一轮的 Key/Value 状态，而不是重新计算。
- 通过掩码替换（Masked Replacement）操作： $K^i = \text{where}(m^i, K^i_{\text{new}}, K^{i-1}_{\text{old}})$ 。
- 优势：该机制不改变注意力算子本身，可无缝集成 Flash Attention 等优化内核，实现真正的推理加速。

2.3 预训练算法

采用两阶段训练策略以稳定优化：

Warm-up 阶段：仅使用交叉熵损失（ $L_{CE}$ ）训练，不启用剪枝正则化。
Pondering 正则化阶段：引入 Ponder Loss ( $L_{ponder}$ $L_{p o n d er}$ )。
- Bottom-K 惩罚：对门控值最小的 $k$ 比例 Token 进行惩罚，鼓励模型在不需要时尽早停止。
- Warmup 调度： $k$ 值（目标剪枝比例）在训练过程中从 0 逐渐增加到目标值（如 0.1），防止模型过早崩溃。

3. 主要贡献 (Key Contributions)

端到端自监督框架：提出了一种在预训练阶段联合学习循环 refinement 和停止行为的框架，无需外部监督信号或人工设定的每 Token/每层剪枝比例。
自动化的 Token 级资源分配：模型能根据 Token 的内在难度自动分配“思考”预算，无需手动调参。
实用的 KV 复用机制：解决了循环模型中训练与推理不一致的问题，使得 Token 级剪枝在实际推理中能够带来显著的速度提升。
大规模验证：在 Pythia 架构（70M 至 2.8B 参数）上进行了广泛的预训练和持续预训练实验，验证了方法的有效性。

4. 实验结果 (Results)

推理效率提升：
- 在 Pythia 70M 和 410M 模型上，AdaPonderLM 在保持与 PonderLM 相当的语言建模困惑度（Perplexity）的同时，减少了约 10% 的推理计算量（FLOPs）（从 4.0x 降至约 3.8x）。
- 在 1.4B 和 2.8B 的持续预训练实验中，同样实现了约 10% 的 FLOPs 降低，且验证损失与基线持平。
下游任务表现：
- 在 LAMBADA, PIQA, WinoGrande, ARC, HellaSwag 等多个下游任务上，AdaPonderLM 的表现与固定深度的 PonderLM 具有竞争力，甚至在 2.8B 规模下零样本（Zero-shot）准确率提升了 2.2%。
- 证明了剪枝冗余计算并未损害生成质量。
自适应行为分析：
- 难度相关性：分析显示，模型倾向于在 Step 2 停止简单 Token（低 NLL），而将高 NLL（困难）Token 保留到后续迭代（Step 3-4），符合自适应计算时间（ACT）原则。
- 策略对比：在等 FLOPs 约束下，学习到的自适应策略优于均匀分布（Uniform）和几何衰减（Geometric）等固定剪枝策略。
消融实验：
- 证明了使用单一共享 MLP 会导致训练不稳定（模型崩溃为全停或全不停），而迭代特定 MLP 是必要的。
- 正则化系数 $\lambda$ 和剪枝比例 $k$ 对性能敏感，较小的值通常能获得更好的损失表现。

5. 意义与结论 (Significance)

效率与质量的平衡：AdaPonderLM 证明了通过自监督学习动态调整推理深度，可以在不牺牲模型能力的前提下显著降低推理成本。
解决“一刀切”问题：打破了传统循环模型固定迭代次数的限制，实现了真正的 Token 级自适应计算，将计算资源精准分配给“困难”的 Token。
工程落地性：通过 KV 复用机制，使得理论上的自适应深度在实际硬件上能够转化为真实的推理加速，为未来高效的大模型推理提供了新的技术路径。
未来方向：虽然目前受限于模型规模，但该方法展示了自监督自适应循环在扩展性上的巨大潜力，未来可探索更大规模模型及结合监督信号进一步优化。

总结：AdaPonderLM 是一种创新的循环语言模型，它通过自监督门控机制和 KV 复用技术，成功实现了 Token 级的自适应推理深度，在降低约 10% 推理计算量的同时，保持了与固定深度模型相当甚至更优的性能，为高效的大模型推理提供了重要的解决方案。