Think-at-Hard: Selective Latent Iterations to Improve Reasoning Language Models
Tianyu Fu, Yichen You, Zekai Chen, Guohao Dai, Huazhong Yang, Yu Wang
提高大型语言模型(LLM)的推理能力,特别是在参数约束下,对于现实世界的应用至关重要。 之前的工作提出了循环变压器,每个令牌分配固定数量的额外迭代,以提高发电质量。 在第一个标准前传之后,而不是口头上,最后一层隐藏状态被反馈为用于改进令牌预测的额外迭代的输入。 然而,我们发现了一个潜在的过度思考现象:在第一次传递之后已经正确的简单令牌预测有时会在额外的迭代中修改为错误。 为了解决这个问题,我们提出了Think-at-Hard(TaH),这是一种动态的潜在思维方法,仅在硬令牌中迭代更深。 它使用轻量级的神经决定器来触发潜在的迭代,仅在标准前传后可能不正确的令牌上触发。 在潜在迭代期间,低等级适应(LoRA)模块将LLM目标从一般的下图预测转变为集中的硬令牌改进。 我们进一步引入了一种双致病注意力机制,将注意力从令牌序列维度扩展到额外的迭代深度维度。 这可以实现跨迭代信息流,同时保持完全顺序的并行。 实验表明,TaH在五个具有挑战性的基准测试中提高了LLM推理性能,同时保持了相同的参数计数。 与对所有输出令牌迭代两次的基线相比,TaH提供了8.1-11.3%的准确率增益,同时从第二次迭代中免除了94%的令牌。 针对使用相同数据进行微调的强单迭代Qwen3型号,它还可提供4.0-5.0%的精度增益。 当允许LoRA和迭代决定器提供不到3%的额外参数时,收益分别增加到8.5-12.6%和5.3-5.4%。 我们的代码可在https://github.com/thu-nics/TaH。
Improving reasoning capabilities of Large Language Models (LLMs), especially under parameter constraints, is crucial for real-world applications. Prior work proposes recurrent transformers, which allocate a fixed number of extra iterations per token to improve generation quality. After the first, standard forward pass, instead of verbalization, last-layer hidden states are fed back as inputs for additional iterations to refine token predictions. Yet we identify a latent overthinking phenomenon: ...