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

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 LLMs 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.

• The paper introduces Test-Time Control, which reframes next-token prediction as a finite-horizon LQR problem so each layer performs explicit planning through a learned value function.
• The hardware-efficient symplectic solver uses parallel matrix operations, fused CUDA kernels, and backward caching to achieve over 10× the throughput of Riccati- and KKT-based baselines while maintaining constant horizon memory.
• TTC-Net improves Sudoku board accuracy to 61.30% and raises Llama-3-7B results to 52.80% on MATH-500, with adjustable test-time horizons providing an additional reasoning-compute tradeoff.

Overview

This paper introduces Test-Time Control (TTC), an architectural mechanism that reframes next-token prediction as solving a finite-horizon Linear-Quadratic Regulator (LQR) problem over latent states at inference time. The authors argue that existing sequence models—attention, linear RNNs, and state-space models—are unified by associative memory: they encode past context and retrieve from it to predict. This memory-centric design corresponds to System 1 cognition, whereas deliberate multi-step planning (System 2) is absent from the forward computation itself. Rather than injecting planning through external reinforcement learning post-training or test-time self-supervised training, TTC internalizes a value function directly into the architecture: each layer solves a receding-horizon optimal control problem conditioned on context, and decodes the first-step optimal action as the next-token representation (2603.09221).

The work makes three coupled contributions: (1) the TTC layer formulation with a differentiable KKT-based backward pass; (2) a hardware-efficient symplectic LQR solver implemented as fused CUDA kernels; and (3) TTC-Net, a hybrid architecture that interleaves TTC layers with attention, evaluated on Sudoku and mathematical reasoning benchmarks.

From test-time training to test-time control

The paper grounds its motivation in the test-time-training view of sequence models, where a memory unit solves an online regression objective over the observed prefix, and prediction amounts to querying the fitted predictor. Attention emerges as a non-parametric Nadaraya-Watson solution; linear RNNs such as DeltaNet emerge as one-step gradient descent on a parametric linear regressor. All of these interpret inference as test-time estimation—fitting a model of the past.

TTC instead interprets inference as test-time decision making. Given a context-encoded initial state $h_0$, the layer synthesizes time-heterogeneous dynamics $\{A_t, B_t\}$ and quadratic costs $\{Q_t, R_t\}$, then solves

$$\min_{u_1,\dots,u_T} \frac{1}{2}\sum_{t=1}^{T}\left(h_t^\top Q_t h_t + u_t^\top R_t u_t\right) \quad \text{s.t.} \quad h_t = A_t h_{t-1} + B_t u_t,$$

and outputs $u_1^*$ as the next-token representation. The value matrices $\{P_t\}$ produced by the Riccati recursion define an explicit value function $V_t(h_t) = -\frac{1}{2}h_t^\top P_t h_t$ within each block, so the model performs planning before prediction. The cost structure accommodates both process rewards (intermediate $Q_t, R_t$) and outcome rewards (terminal $Q_T$).

For end-to-end learning, the authors derive gradients via implicit differentiation of the KKT system: backpropagation requires solving a second ("dual") LQR with zero initial state and an affine term involving $\nabla_o \ell$. Gradients with respect to $\{A_t, B_t\}$0 are outer products of primal and dual states, actions, and co-states. Training thus becomes a nested process—an inner loop solving the control problem per token, an outer loop updating world-model parameters—which the authors frame as inverse RL when fine-tuning on expert demonstrations.

Hardware co-design: symplectic iteration

Classical Riccati solvers are both compute- and I/O-bound: they require $\{A_t, B_t\}$1 sequential matrix inversions poorly suited to accelerators. The central algorithmic result replaces the Riccati recursion with a symplectic iteration: exploiting the fact that the state–co-state transition matrix $\{A_t, B_t\}$2 is symplectic, the optimal first-step action can be recovered from a reverse cumulative product of symplectic factors $\{A_t, B_t\}$3, requiring only one dense inversion ($\{A_t, B_t\}$4) outside the product. The per-step inversions inside $\{A_t, B_t\}$5 are independent across time and fully parallelizable; all remaining sequential work is dense matrix multiplication amenable to Tensor Cores.

Three further optimizations make this practical:

• Structured parameterization: diagonal $\{A_t, B_t\}$6 and $\{A_t, B_t\}$7 reduce nontrivial dense inversions from $\{A_t, B_t\}$8 to $\{A_t, B_t\}$9, without limiting expressivity relative to dense SSMs.
• Kernel fusion: $\{Q_t, R_t\}$0 admits a three-factor decomposition, so factors stream into SRAM and cumulative products run block-wise in a single fused kernel, with row-wise normalization preventing overflow from reciprocal-pair eigenvalue growth inherent to symplectic products.
• Backward caching: because the dual LQR differs from the primal only in its final step, the LU factorization of $\{Q_t, R_t\}$1 and part of the cumulative product computed in the forward pass are reused, eliminating a second reverse iteration during backpropagation.

Benchmarks on an H200 show over 10× higher throughput than Riccati- and KKT-based baselines, with constant memory scaling in horizon where baselines hit out-of-memory errors. This result is what makes embedding LQR into thousands of tokens' worth of layers feasible; without it, the architectural proposal would be computationally impractical.

TTC-Net architecture and training

TTC-Net inserts a TTC layer between attention and MLP every 8 transformer blocks. Each token's attention output is projected to an initial state $\{Q_t, R_t\}$2; contextualization generates time-modulation coefficients $\{Q_t, R_t\}$3 that discount parameters across the horizon, yielding time-heterogeneous dynamics with guaranteed positive semi-definiteness via Cholesky reparameterization of basis matrices. A multi-head structure partitions the state into blocks of size 16 with shared bases. Output projections are zero-initialized so fine-tuning from a pretrained checkpoint starts identical to the backbone.

Two training choices matter for horizon generalization. Mixed-horizon training samples horizons from a truncated Poisson log-normal distribution (mean 8, max 32), avoiding distribution shift when the test-time horizon differs. At inference, the horizon $\{Q_t, R_t\}$4 is freely adjustable, exposing a native test-time scaling axis: longer horizons cost more FLOPs but improve accuracy, and the model generalizes beyond the training maximum to $\{Q_t, R_t\}$5.

Empirical results

Sudoku. On 9×9 boards with 17–34 givens, a 32-layer TTC-Net outperforms Transformer, Mamba, Mamba2, GDN, and Samba under matched capacity: 61.30% vs. 58.50% board accuracy for single-step completion (+2.8% over the runner-up), and 93.40% vs. 90.40% for iterative multi-step completion. Cell-level accuracy reaches 97.33%. The gain is attributed to the TTC objective acting as a smooth surrogate for constraint satisfaction, with digit placement as linear transitions.

Mathematical reasoning. Fine-tuning Llama-3-Instruct-7B with TTC adapters on OpenThoughts2-114K plus 800K curated reasoning examples yields consistent gains over hybrid adapters built from attention, RetNet, Mamba, GDN, and MesaNet:

Model	MATH-500	AMC Pass@8	AIME24 Pass@8	AIME25 Pass@8
Base model	25.00	31.32	0.00	0.00
Full finetuning	46.80	46.98	6.67	0.00
+ GDN (best adapter baseline)	47.80	37.35	3.33	6.67
TTC-Net	52.80	54.22	20.00	20.00

The headline claims are up to +27.8% absolute improvement on MATH-500 over the base model and 2–3× Pass@8 gains on AMC and AIME, including nonzero performance on AIME 2025 where every baseline scores zero. Notably, memory-layer adapters generally fail to stably exceed plain SFT of the base model, while TTC-Net does—a contrast the authors use to argue that the control-based objective provides an inductive mechanism absent from additional memorization capacity. The large Pass@8 (rather than Avg@8) gains suggest TTC extends the effective capability boundary of the base model rather than merely sharpening existing behavior, which aligns with prior observations that post-training alone cannot lift intrinsic ceilings.

Ablations. Time-homogeneous parameterizations degrade both absolute accuracy and horizon generalization; fixed-horizon training collapses when tested at a different horizon (31.50% at $\{Q_t, R_t\}$6 vs. 53.60% for the full model); uniform horizon sampling matches PLN sampling but roughly doubles average training compute; and distributing TTC layers uniformly (8:1 interleaving) beats stacking them consecutively (16:2). Increasing TTC density improves accuracy monotonically but less cost-effectively than extending the test-time horizon.

Limitations and open questions

The paper is candid about several constraints. Theoretically, while a single TTC layer has a well-defined optimization semantics, how multiple TTC layers interact and jointly represent dynamics within a deep transformer remains unanalyzed. Empirically, the parameterization restricts dynamics to structured linear systems with diagonal $\{Q_t, R_t\}$7 and $\{Q_t, R_t\}$8; more expressive (including nonlinear) formulations compatible with hardware constraints are left open. Evaluation covers a 7B model with adapter-style insertion; comprehensive assessment across larger scales and all training stages (pretraining through RL) is not provided. The comparison against TTD—a concurrent model-free, slow-weight test-time RL method—is qualitative only, and whether fast-weight value-based planning composes with slow-weight policy adaptation is untested. Finally, the claim that control objectives unlock capabilities beyond the base model's ceiling rests on Pass@8 evidence at limited sample counts on AIME-sized benchmarks, where variance is high.

Conclusion

This paper positions optimal control as an architectural primitive for LLM reasoning: TTC layers solve contextualized finite-horizon LQR problems in the forward pass, carry an explicit value function, and differentiate exactly through KKT conditions. The symplectic solver—with parallelizable inversions, kernel fusion, and backward caching—delivers the >10× throughput needed to deploy this at LLM scale, and the resulting hybrid model shows consistent reasoning gains over memory-only counterparts, with a native test-time scaling axis via planning horizon. The work reframes test-time adaptation as structured decision making rather than estimation or parameter updates, though questions about deep-stack composition, richer dynamics, and scale remain open.