テスト時学習を超えて：ハードウェア効率の良い最適制御による推論学習

要旨

連想記憶は長らく系列モデルの設計を支えてきた。しかし、人間の推論は想起を超え、未来状態を予測し目標指向の行動を選択するものであり、これは現代の言語モデルが強く求めながらも本質的に備えていない能力である。従来研究では強化学習やテスト時訓練が用いられてきたが、計画機能はモデル構造の外部に留まっていた。本研究では推論を最適制御問題として定式化し、テスト時制御（TTC）層を提案する。TTC層は推論時に潜在状態に対して有限時間LQR計画を実行し、ニューラル構造内に価値関数を表現し、これを予測前の計画を可能にする入れ子目的関数として活用する。拡張性を確保するため、シンプレクティック定式化に基づくハードウェア効率の良いLQRソルバーを導出し、オーバーヘッドを最小化した融合CUDAカーネルとして実装することで並列実行を可能にした。事前学習済み大規模言語モデルにアダプタとして統合されたTTC層は、MATH-500で最大27.8%の数学推論性能向上、AMCとAIMEではPass@8で2～3倍の改善を示し、最適制御を構造的構成要素として組み込むことが、テスト時訓練を超えた推論のための効果的かつ拡張性の高いメカニズムを提供することを実証した。

English

Associative memory has long underpinned the design of sequential models. Beyond recall, humans reason by projecting future states and selecting goal-directed actions, a capability that modern language models increasingly require but do not natively encode. While prior work uses reinforcement learning or test-time training, planning remains external to the model architecture. We formulate reasoning as optimal control and introduce the Test-Time Control (TTC) layer, which performs finite-horizon LQR planning over latent states at inference time, represents a value function within neural architectures, and leverages it as the nested objective to enable planning before prediction. To ensure scalability, we derive a hardware-efficient LQR solver based on a symplectic formulation and implement it as a fused CUDA kernel, enabling parallel execution with minimal overhead. Integrated as an adapter into pretrained LLMs, TTC layers improve mathematical reasoning performance by up to +27.8% on MATH-500 and 2-3x Pass@8 improvements on AMC and AIME, demonstrating that embedding optimal control as an architectural component provides an effective and scalable mechanism for reasoning beyond test-time training.

テスト時学習を超えて：ハードウェア効率の良い最適制御による推論学習

Beyond Test-Time Training: Learning to Reason via Hardware-Efficient Optimal Control

要旨

Support