Think-at-Hard：選択的潜在反復による推論言語モデルの改善

要旨

大規模言語モデル（LLM）の推論能力向上、特にパラメータ制約下での性能向上は実世界応用において重要である。従来研究では、生成品質を改善するためにトークンごとに固定数の追加反復を割り当てるリカレントトランスフォーマーが提案されている。最初の標準的な順伝搬処理の後、言語化を行う代わりに最終層の隠れ状態を入力としてフィードバックし、追加反復によってトークン予測を洗練させる。しかし我々は潜在的な「過剰思考（overthinking）」現象を特定した：最初の処理で既に正解している容易なトークン予測が、追加反復において誤りに修正される場合があるのである。この問題に対処するため、我々は困難なトークンに対してのみ深く反復を行う動的潜在思考手法「Think-at-Hard（TaH）」を提案する。本手法は軽量なニューラルデサイダーを用いて、標準的な順伝搬処理後に誤りである可能性が高いトークンのみで潜在反復をトリガーする。潜在反復中は、LoRA（Low-Rank Adaptation）モジュールによってLLMの目的を一般的な次トークン予測から焦点を絞った困難トークンの洗練へと移行させる。さらに我々は、注意力機構をトークン系列次元から反復深度次元へ拡張するデュオ因果的注意力機構を導入する。これにより完全な逐次並列性を維持しつつ、反復間の情報流れを可能にする。実験結果では、TaHが同じパラメータ数を維持しながら、5つの困難なベンチマークでLLMの推論性能を向上させることを示す。全ての出力トークンに対して2回反復するベースラインと比較して、TaHは94%のトークンを第2反復から免除しつつ8.1-11.3%の精度向上を達成した。同じデータでファインチューニングされた強力な単一反復型Qwen3モデルに対しても、4.0-5.0%の精度向上を実現した。LoRAと反復デサイダーによる3%未満の追加パラメータを許容した場合、精度向上はそれぞれ8.5-12.6%および5.3-5.4%に増加した。実装コードはhttps://github.com/thu-nics/TaHで公開している。

English

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: easy token predictions that are already correct after the first pass are sometimes revised into errors in additional iterations. To address this, we propose Think-at-Hard (TaH), a dynamic latent thinking method that iterates deeper only at hard tokens. It employs a lightweight neural decider to trigger latent iterations only at tokens that are likely incorrect after the standard forward pass. During latent iterations, Low-Rank Adaptation (LoRA) modules shift the LLM objective from general next-token prediction to focused hard-token refinement. We further introduce a duo-causal attention mechanism that extends attention from the token sequence dimension to an additional iteration depth dimension. This enables cross-iteration information flow while maintaining full sequential parallelism. Experiments show that TaH boosts LLM reasoning performance across five challenging benchmarks while maintaining the same parameter count. Compared with baselines that iterate twice for all output tokens, TaH delivers 8.1-11.3% accuracy gains while exempting 94% of tokens from the second iteration. Against strong single-iteration Qwen3 models finetuned with the same data, it also delivers 4.0-5.0% accuracy gains. When allowing less than 3% additional parameters from LoRA and the iteration decider, the gains increase to 8.5-12.6% and 5.3-5.4%, respectively. Our code is available at https://github.com/thu-nics/TaH.

Think-at-Hard：選択的潜在反復による推論言語モデルの改善

Think-at-Hard: Selective Latent Iterations to Improve Reasoning Language Models

要旨

Support