Beyond Test-Time Training: Learning to Reason via Hardware-Efficient Optimal Control
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.
69b1acf2206bad330cd7fe58
Peihao Wang
69b1acf2206bad330cd7fe59
Shan Yang
69b1acf2206bad330cd7fe5a
Xijun Wang
69b1acf2206bad330cd7fe5b
Tesi Xiao
69b1acf2206bad330cd7fe5c
Xin Liu
69b1acf2206bad330cd7fe5d
Changlong Yu
69b1acf2206bad330cd7fe5e
Yu Lou
69b1acf2206bad330cd7fe5f
Pan Li
69b1acf2206bad330cd7fe60
Zhangyang Wang
69b1acf2206bad330cd7fe61
Ming Lin
69b1acf2206bad330cd7fe62
René Vidal
Beyond Test-Time Training: Learning to Reason via Hardware-Efficient Optimal Control
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.
Test-Time Control layer integrates optimal control theory into language models for enhanced reasoning, using LQR planning and hardware-efficient solvers to improve mathematical problem-solving performance.
associative memory sequential models reinforcement learning test-time training optimal control Test-Time Control (TTC) layer finite-horizon LQR planning latent states value function nested objective hardware-efficient LQR solver symplectic formulation fused CUDA kernel pretrained LLMs mathematical reasoning MATH-500 AMC AIME