arXiv:2511.08577cs.CLcs.AI2025-11中稿 · ICML

让大模型只在关键位置反复思考,提升推理准确率同时大幅减少计算。

Think-at-Hard: Dynamic Looped Transformers for Improved Reasoning

  • 用轻量决策器判断哪些词需要反复思考,其他直接跳过。
  • 在9个任务上比固定循环模型高3.8%-4.4%,跳过93%的冗余迭代。
  • 适合追求高效精准推理的场景,如代码生成和数学题求解。

提升大语言模型的推理能力,尤其是在参数受限的情况下,对实际应用至关重要。环形变换器通过多次潜空间迭代来优化每个标记,超越单次前向传播。然而,我们发现存在潜在的过度思考现象:多数标记在首次预测后已正确,但后续迭代中可能被错误修正。我们探究是否可选择性跳过潜空间迭代以提高准确性,发现理想策略可使性能提升高达7.3%。受此启发,我们提出Think-at-Hard(TaH),一种专为选择性迭代优化的环形变换器。TaH采用轻量神经决策器,在标准前向传播后判断可能出错的标记才触发潜空间迭代。迭代过程中,深度感知的低秩适配(LoRA)模块将目标从通用下一个标记预测转向聚焦于困难标记的精细修正。双因果注意力机制将注意力扩展至迭代深度维度,实现跨迭代信息流动的同时保持序列并行性。在九个基准测试上的实验显示,相同参数量下,TaH比始终迭代的基线模型提升3.8%-4.4%,且跳过93%的迭代;当允许增加<3%参数用于LoRA与决策器时,相比单次迭代Qwen3基线,性能进一步提升6.1%-6.8%。代码已开源。

原文摘要 · Abstract (English)

Improving the reasoning abilities of Large Language Models (LLMs), especially under parameter constraints, is crucial for real-world applications. Looped transformers address this by performing multiple latent iterations to refine each token beyond a single forward pass. However, we identify a latent overthinking phenomenon: most token predictions are already correct after the first pass, but are sometimes revised into errors in later iterations. We ask whether selectively skipping latent iterations can improve accuracy, and reveal significant potential with an oracle iteration policy that boosts performance by up to 7.3%. Motivated by this, we propose Think-at-Hard (TaH), a looped transformer optimized for selective iteration. TaH employs a lightweight neural decider to trigger latent iteration, only at tokens likely to be incorrect after the standard forward pass. During latent iterations, depth-aware Low-Rank Adaptation (LoRA) modules shift the objective from general next-token prediction to focused hard-token refinement. A duo-causal attention mechanism extends attention from the token sequence dimension to an additional iteration depth dimension, enabling cross-iteration information flow with full sequential parallelism. Experiments on nine benchmarks show consistent gains across math, QA, and coding tasks. With identical parameter counts, TaH outperforms always-iterate baselines by 3.8-4.4% while skipping iterations on 93% of tokens, and exceeds single-iteration Qwen3 baselines by 3.0-3.8%. When allowing <3% more parameters from LoRA and decider, the gains further increase to 5.3-6.2% and 6.1-6.8%, respectively. Our code is available at https://github.com/thu-nics/TaH.

推理增强环形结构动态计算效率优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。