Command Palette
Search for a command to run...
Tianyu Fu Yichen You Zekai Chen Guohao Dai Huazhong Yang Yu Wang

摘要
在参数受限条件下,提升大语言模型(LLM)的推理能力对于实际应用至关重要。以往研究提出循环Transformer(recurrent transformers)架构,通过为每个 token 固定增加若干次额外迭代来提升生成质量。在完成首次标准前向传播后,不再直接输出结果,而是将最后一层隐藏状态作为输入,进行额外迭代以优化 token 预测。然而,我们发现存在一种隐含的“过度思考”现象:在首次前向传播中已正确预测的简单 token,可能在后续迭代中被错误修正。为解决该问题,我们提出 Think-at-Hard(TaH)——一种动态的隐式思考机制,仅在难以预测的 token 上进行深度迭代。该方法引入一个轻量级神经决策器,仅在标准前向传播后判断可能出错的 token 上触发隐式迭代。在隐式迭代过程中,通过低秩适配(Low-Rank Adaptation, LoRA)模块,将 LLM 的目标从通用的下一个 token 预测,转变为聚焦于难例 token 的精细化修正。此外,我们设计了一种双因果注意力机制(duo-causal attention),将注意力范围从传统的 token 序列维度拓展至额外的迭代深度维度,实现跨迭代的信息流动,同时保持完整的序列并行性。实验结果表明,TaH 在五个具有挑战性的推理基准上显著提升了 LLM 的性能,且参数量保持不变。相比对所有输出 token 均进行两次迭代的基线方法,TaH 在免除 94% token 进行第二次迭代的同时,实现了 8.1%–11.3% 的准确率提升。相较于使用相同数据微调的强基线单次迭代 Qwen3 模型,TaH 仍带来 4.0%–5.0% 的准确率增益。当允许引入少于 3% 的额外参数(来自 LoRA 与迭代决策器)时,准确率提升进一步增至 8.5%–12.6% 与 5.3%–5.4%。相关代码已开源,地址为:https://github.com/thu-nics/TaH。