Beyond Test-Time Memory: State-Space Optimal Control for LLM Reasoning

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-3$\times$ 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.

Large language models are increasingly expected to solve problems that require deliberate reasoning, such as math, coding, and logical puzzles. Reinforcement learning has helped make these models more goal-directed, but it is usually applied as an external training procedure: it changes the model’s behavior, while the model architecture itself still largely predicts by recalling and recombining patterns from past context. This creates a gap between the reasoning behavior we want and the memory-based computation that most language models natively perform. We introduce Test-Time Control, a new layer that brings planning directly into the model's forward pass. Instead of only retrieving useful information from context, this layer considers possible future internal states and chooses the next step that best supports the overall goal. The planning problem is based on a classical control method from engineering, and we redesign its solver so it can run efficiently on modern GPUs at scale. In experiments on logical and mathematical reasoning tasks, models with this layer solve more problems than comparable models that rely only on memory-based components or fine-tuning. These results suggest that reasoning can be strengthened not only by external training objectives, but also by changing the architecture so that the model can plan before it answers. More broadly, this work offers a step toward language models that combine memory-based prediction with built-in goal-directed reasoning.