GradCuit: クレジット割り当て勾配フローによる堅牢かつ解釈可能なテスト時潜在推論

要旨

最適化ベースの潜在推論は、モデルパラメータを凍結したまま、テスト時にインスタンス固有の連続状態を最適化することにより、大規模言語モデルの出力を改善する。しかし、既存手法は通常、これらの状態をデコードされたトークンを介して推論軌跡に接続するため、シーケンスレベルのクレジット割り当てが間接的となり、潜在状態の更新がその後の推論をどのように形成するかが不明瞭である。我々は、選択したTransformer層において、プロンプトの隠れ表現と生成された継続部分の間に最適化可能な潜在状態を挿入するGradCuit（勾配を回路に通す）を導入する。因果的自己注意により、各継続トークンの対数確率は、残りのTransformerブロックを通じて先行するすべての潜在状態への微分可能な経路を持ち、継続全体からの報酬重み付き勾配を潜在状態に直接割り当てることが可能になる。5つの指示チューニング済みバックボーン、3つの推論ベンチマーク、2つの解答形式にわたって、GradCuitは平均精度64.5%を達成し、思考連鎖プロンプティングを6.6パーセントポイント、最強の競合手法を2.4ポイント上回る。GradCuitはまた、より高い堅牢性を示す：7つの学習率設定にわたってLatentSeekを一貫して上回りながら、精度の標準偏差を1.53から0.82に低減し、そのランダムウォーク変種でさえLatentSeekと競争力を維持する。解釈可能性については、トークンレベルの勾配帰属により、潜在状態の影響が推論接続トークンに集中することが明らかになり、層分析では初期から中間のTransformer層が最も効果的な最適化空間であることが特定される。結果フィードバックから内部推論を直接最適化することにより、GradCuitは、LLMが出力を単に再生成、サンプリング、再ランキングするのではなく、推論の仕方そのものを適応させる、堅牢かつ解釈可能なテスト時スケーリングの新たな軸を切り開く。

English

Optimization-based latent reasoning improves large language model outputs by optimizing instance-specific continuous states at test time while keeping model parameters frozen. Existing methods, however, typically connect these states to the reasoning trajectory through decoded tokens, making sequence-level credit assignment indirect and obscuring how latent updates shape subsequent reasoning. We introduce GradCuit (gradient through circuit), which inserts optimizable latent states at a selected Transformer layer between the hidden representations of the prompt and the generated continuation. Causal self-attention provides every continuation-token log-probability with a differentiable path to every preceding latent state through the remaining Transformer blocks, enabling reward-weighted gradients from the entire continuation to be assigned directly to the latents. Across five instruction-tuned backbones, three reasoning benchmarks, and two answer formats, GradCuit achieves an average accuracy of 64.5%, outperforming chain-of-thought prompting by 6.6 percentage points and the strongest competing method by 2.4 points. GradCuit also demonstrates greater robustness: across seven learning-rate settings, it consistently outperforms LatentSeek while reducing the standard deviation of accuracy from 1.53 to 0.82, and even its random-walk variant remains competitive with LatentSeek. For interpretability, token-level gradient attribution reveals that latent influence concentrates on reasoning-connector tokens, while layer analysis identifies early-to-middle Transformer layers as the most effective optimization space. By directly optimizing internal reasoning from outcome feedback, GradCuit opens a new axis of robust and interpretable test-time scaling, where LLMs adapt how they reason rather than merely regenerate, sample, or rerank outputs.