清华TaH模型：跳过93%无效迭代，准确率反升的深度解析

在o1、R1等“推理模型”的推动下，“增加模型思考时间”已成为处理复杂任务的主流策略。无论是扩展思维链、增加测试时计算量，还是引入更深层的内部推理循环，其本质都是通过投入更多计算资源来换取答案的可靠性。

但一个核心问题常被忽略：模型是否真的需要在生成每一个词元（token）时，都进行额外的“思考”？

这对于参数受限的小型模型至关重要。小模型成本低、速度快、适合边缘部署，但在数学、代码等复杂推理任务中，往往因少数关键token预测失误而全盘皆输。现有的“循环Transformer”方案试图让模型在生成每个token前，将隐藏状态送回进行“潜空间迭代”，以此在不增加参数的前提下提升计算深度。

然而，清华大学、无问芯穹及上海交通大学的研究团队发现，这一过程存在效率瓶颈。他们的研究表明，许多token在首次前向传播时已预测正确，后续的潜空间迭代反而可能“矫枉过正”，将原本正确的答案改错。研究团队将这一现象定义为“潜空间过度思考”。

基于此洞察，他们提出了Think-at-Hard（TaH）方法。这是一种面向小模型的选择性潜空间迭代框架，其核心是让模型学会“仅在真正困难的token上投入额外算力”。通过轻量级后训练，可将标准模型转化为仅在关键位置进行循环计算的优化版本。该工作已入选ICLR LIT Workshop最佳论文候选，并被ICML 2026接收。

简言之，潜空间迭代是一把双刃剑：它能修正错误，也可能引入错误。

核心发现：对简单token进行迭代反而损害性能

已有研究指出，语言模型推理过程中，并非所有token都同等重要。决定推理路径走向的，往往是那些处于逻辑转折、因果连接或代表中间结论的关键位置。

为量化“选择性迭代”的潜力，团队设计了一个“先知策略”：仅当模型首次预测错误时，才允许其进行潜空间迭代；若首次预测正确，则直接输出。实验表明，仅凭这一理想化策略，模型在下游任务上的性能最高可提升7.3%，且仅需对11%到19%的token进行二次迭代。若采用为选择性迭代优化的TaH架构，性能提升甚至可超过25%。

这清晰地证明，推理计算资源的分配应细化至token级别。难题中包含简单token，简单题中也可能出现关键token。更重要的是，对简单token强行进行额外计算不仅是资源浪费，更会导致一部分原本正确的预测在二次迭代中被改错，这正是“潜空间过度思考”的体现。

TaH：在困难处驻足深思

TaH的思路直接而高效：让简单token快速输出，仅让困难token进入下一轮迭代思考。

具体实现上，TaH在模型中引入了一个轻量级“迭代决策器”。每完成一轮潜空间迭代后，该决策器会基于骨干网络的状态，预测一个“继续迭代”的概率。若概率低于阈值，模型直接输出当前token；若高于阈值，则进入下一轮迭代。

在实际推理中，TaH平均每个token仅执行约1.07次迭代，相当于跳过了约93%的token的二次迭代计算。相比“所有token都想两遍”的均匀策略，TaH成功将宝贵算力集中到了更易出错、且更能影响全局推理方向的关键位置上。

为使这种动态深度策略切实提升精度与效率，TaH在模型架构和训练策略上均进行了针对性设计。

Duo-Causal Attention架构

选择性迭代将模型处理的序列结构，从一维的token序列转变为“token位置 × 迭代深度”的二维网格。TaH将大模型序列维度的因果注意力机制，扩展至token维度和迭代深度的二维平面。如图所示，对于第i个token的第d次迭代，其查询可关注所有前序位置中、迭代深度不超过d的键和值。这既允许了跨迭代深度的信息流动，又保留了训练时序列维度计算的全并行性。

Depth-Aware LoRA架构

研究团队观察到，模型的首次迭代主要负责常规的下一个token预测，而更深层的迭代则专注于修正当前的困难token。因此，TaH仅在深度大于1的迭代中启用LoRA适配器，让LoRA专注于学习困难token的修正方向。配合跨迭代的残差连接，深层迭代被自然地学习为“在前一轮预测基础上进行修正”，而非从头开始预测。

两阶段训练策略

这里存在一个耦合难题：决策器的判断依赖于骨干网络的预测质量，而骨干网络的训练目标又依赖于决策器决定的迭代深度，导致端到端训练极不稳定。TaH采用解耦的两阶段方案：首先，使用静态的“先知策略”训练骨干模型；随后，冻结骨干模型，单独训练决策器以模仿“先知”的继续或停止决策。该方法显著提升了训练的稳定性和收敛速度。

结果：更少迭代，更强推理

论文在Qwen3的0.6B、1.7B和4B三个规模的骨干模型上验证了TaH。训练数据来自Open-R1中数学、问答和代码任务的均衡混合，并在GSM8K、MATH500、AMC23、AIME25、OlympiadBench、GPQA-Diamond、MMLU-STEM、HumanEval++和MBPP++共9个基准上进行了全面评测。所有对比方法均使用完全相同的训练数据，在相同的预训练骨干模型上进行后训练。

准确性显著提升。 在0.6B和1.7B的严格参数限制下，TaH相比标准Qwen3模型提升了3.0%到3.8%。TaH+在增加不超过3%额外参数（来自决策器等组件）的情况下，将提升幅度进一步扩大至5.3%到6.2%。与同类的循环Transformer方法Ouro相比，TaH取得了3.8%到4.4%的优势，而TaH+的优势达到了6.1%到6.8%。

计算效率更高。 TaH平均每个token仅执行1.07次迭代，完成问答的平均FLOPs和显存访问量相比标准模型仅增加4%到5%。在真实解码测试中，TaH相比“始终迭代”策略，显存占用降低了1.48倍，解码速度快了2.48倍，同时准确率更高。

迭代选择具有可解释的语义偏好。 一个有趣的现象是，TaH自动学习到了带有明显语义偏好的迭代行为。在验证集上，“But”和“So”是最常触发额外迭代的token，概率分别达到34%和18%。这些词大多对应着转折、因果和推理方向的切换，正是复杂推理中最可能决定后续路径的关键位置。

消融实验

为验证TaH中每项设计的必要性，研究团队进行了系统的消融实验。

模型架构方面。 将动态迭代深度的决策器替换为固定深度策略（Always-1或Always-2），基准测试性能平均分别下降6.1%和16.4%，这证明了选择性迭代本身优于统一深度。将duo-causal注意力机制替换为不同形式的传统因果注意力，测试性能下降5.4%到8.5%，证实了跨迭代深度信息流动的重要性。移除depth-aware LoRA与跨迭代残差连接，效果下降4.9%，确认了在辅助不同迭代进行优化目标转变时，专用架构的关键作用。

训练策略方面。 相比TaH的两阶段训练，简单地监督所有深度的预测会使基准测试性能下降4.3%，说明不同迭代层应承担不同的优化目标。将训练时的迭代策略从静态“先知”换成基于决策器的或动态的“先知”策略，会因骨干网络与决策器的耦合而导致训练不稳定甚至崩溃，这证明了TaH两阶段训练的必要性。

总结与展望

TaH的意义不仅在于提出了一种新的循环Transformer变体与后训练方法，更在于它探索了如何将“测试时缩放”推向更细粒度的token级别。TaH表明，更智能的动态算力分配，甚至可以比单纯使用更高算力的均匀计算模型带来更好的效果。这为后续在有限资源下最大化模型推理能力的研究，提供了新的思路与启发。

参考文献

[1] Jaech, A., Kalai, A., Lerer, A., et al. OpenAI o1 system card. arXiv preprint arXiv:2412.16720, 2024.

[2] Guo, D., Yang, D., Zhang, H., et al. DeepSeek-R1: Incentivizing reasoning capability in LLMs via reinforcement learning. arXiv preprint arXiv:2501.12948, 2025.

[3] Yang, A., Li, A., Yang, B., et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025.

[4] Abdin, M., Aneja, J., Awadalla, H., et al. Phi-3 technical report: A highly capable language model locally on your phone. arXiv preprint arXiv:2404.14219, 2024.

[5] Team, M., Xiao, C., Li, Y., et al. MiniCPM4: Ultra-efficient LLMs on end devices. arXiv preprint arXiv:2506.07900, 2025.

[6] Hutchins, D., Schlag, I., Wu, Y., Dyer, E., and Neyshabur, B. Block-recurrent transformers. Advances in Neural Information Processing Systems, 35:33248–33261, 2022.

[7] Saunshi, N., Dikkala, N., Li, Z., Kumar, S., and Reddi, S. J. Reasoning with latent thoughts: On the power of looped transformers. arXiv preprint arXiv:2502.17416, 2025.

[8] Zhu, R.-J., Wang, Z., Hua, K., et al. Scaling latent reasoning via looped language models. arXiv preprint arXiv:2510.25741, 2025.

[9] Wu, Y., Wang, Y., Ye, Z., Du, T., Jegelka, S., and Wang, Y. When more is less: Understanding chain-of-thought length in LLMs. arXiv preprint arXiv:2502.07266, 2025.

[10] Wang, S., Yu, L., Gao, C., et al. Beyond the 80/20 rule: High-entropy minority tokens drive effective reinforcement learning for LLM reasoning. arXiv preprint arXiv:2506.01939, 2025.

[11] Fu, T., Ge, Y., You, Y., et al. R2R: Efficiently na vigating divergent reasoning paths with small-large model token routing. arXiv preprint arXiv:2505.21600, 2025.

[12] Hu, E. J., Shen, Y., Wallis, P., et al. LoRA: Low-rank adaptation of large language models. ICLR, 2022.

[13] Hugging Face. Open R1: A fully open reproduction of DeepSeek-R1, January 2025. URL https://github.com/huggingface/open-r1.