[論文レビュー] Beyond Test-Time Training: Learning to Reason via Hardware-Efficient Optimal Control

Peihao Wang, Shan Yang+9 more|arXiv (Cornell University)|Mar 10, 2026

Reinforcement Learning in Robotics被引用数 0

ひとこと要約

論文は最適制御 TTC 層を事前学習済み LLM に組み込み、ハードウェアに優しい LQR ソルバーを用いた推論時計画を可能にすることで、数学・記号タスクの推論を改善します。

ABSTRACT

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.

研究の動機と目的

• 記憶ベースの予測を超える推論の動機付けをモデルアーキテクチャに組み込む。
• 推論時に潜在状態上で有限ホライズンの LQR 計画を行う Test-Time Control (TTC) 層を導入する。
• 制御層を通じてエンドツーエンド学習と勾配伝搬を可能にする微分可能な TTC フレームワークを開発する。
• 推論時の高スループット計画のため CUDA に統合されたハードウェア効率の高い対称性 LQR ソルバーを設計する。
• 難易度の高い推論ベンチマークで事前学習済み LLM にアダプターとして TTC 層を統合した場合の利得を実証する。

提案手法

• 線形ダイナミクスと2次コストを持つ receding-horizon LQR を用いた最適制御問題として推論を定式化する。
• TTC(h0, A_t,B_t,Q_t,R_t) を定義し、コンテキスト潜在状態を第一ステップの最適制御作用へ写像する。
• TTC 層を経由して勾配を逆伝播できる微分可能な KKT ベースの定式化（内側の LQR と外側の世界モデル更新）を導出する。
• 並列行列積と1つの密結合逆行列を用いたハードウェア効率の高い対称性反復ソルバーを開発する。
• A_t と R_t を対角行列としてパラメータ化し定数時間の反転を可能にし、CUDA カーネルに対称性反復を組み合わせてスループットを向上させる。
• TTC をメモリベースモジュールと交互に組み合わせたアダプターとして TTC-Net アーキテクチャを形成する。

実験結果

リサーチクエスチョン

• RQ1 言語モデル内に推論を最適制御問題として内在化することで推論は実現できるか。
• RQ2 テスト時制御層を大規模言語モデル内で微分可能かつエンドツーエンドで訓練可能にするにはどうすればよいか。
• RQ3 推論時のスケーラブルで低オーバーヘッドな計画を実現するハードウェアとアルゴリズムの協創設計戦略は何か。
• RQ4 TTC 強化モデルは数学的・記号的推論ベンチマークで記憶のみの構成より優れているか。

主な発見

• TTC 層は MATH-500 における数学的推論性能を最大で 27.8% 向上させる。
• TTC を事前学習済み LLM に組み込むと AMC および AIME で Pass@8 の改善が 2–3 倍となる。
• 対称性 LQR ソルバーにより並列化可能なハードウェア効率の高い計画を最小オーバーヘッドで実現できる。
• TTC-Net は難易度の高い推論タスクで一貫して記憶ベースモデルを上回る。
• 構造化対角パラメータ化と統合 CUDA カーネルにより推論時および訓練時の計画の高いスループットを実現する。
• このアプローチは、ハードウェア効率の高い最適制御層を介して訓練と推論を統合する統一的な視点を提供する。