PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Peihao Wang, Shan Yang, Xijun Wang, Tesi Xiao, Xin Liu, Changlong Yu, Yu Lou, Pan Li, Zhangyang Wang, Ming Lin, René Vidal

arXiv:2603.09221v2cs.LG

TL;DR

Current LLM architectures lack a dedicated mechanism for System 2-style planning, while long-context inference makes conventional LQR solvers difficult to apply. The paper introduces TTC layers that plan latent-state trajectories through optimal control and integrates them into pretrained LLMs as TTC-Net; the supplied passages report ablation evidence about horizon sampling and layer interleaving, alongside open theoretical and evaluation questions.

  • Problem

    Current LLM architectures largely rely on next-token prediction from past context and lack a dedicated architectural mechanism for deliberate, multi-step planning.

  • Method

    TTC layers formulate reasoning as optimal control over latent states, and TTC-Net interleaves them with attention and MLP modules as adapters in pretrained LLMs.

  • Results

    Ablation studies compare time modulation, horizon sampling, and TTC-layer insertion intervals while testing generalization across inference-time horizons.

  • Takeaways & Limitations

    Embedding planning into the model architecture provides a structured decision-making framework that combines memory with long-horizon planning.

  • Takeaways & Limitations

    The joint interaction of multiple TTC layers remains theoretically unclear, and broader evaluation across larger-scale models and all training stages remains open.

Abstract

from arXiv · show

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.

1 Introduction

The paper argues that memory-centric language models lack dedicated mechanisms for deliberate planning and introduces TTC-Net to internalize optimal-control reasoning before prediction. TTC layers encode value functions, support differentiable and hardware-efficient LQR planning, and improve challenging reasoning benchmarks.

  • Motivation: Current sequence models primarily generate tokens through associative memory, while deliberate multi-step planning remains architecturally unsupported.The paper frames this gap as a difference between System 1-style retrieval and System 2-style reasoning.
  • Motivation: Reinforcement learning makes models more goal-directed but is generally external to pretraining and decoupled from the core inference mechanism.The paper presents this separation as a limitation of relying on RL alone to overcome memory-based reasoning ceilings.
  • Approach: TTC layers support value functions, context-dependent world modeling, differentiable nested optimization, and time-heterogeneous dynamics and costs.The inner loop solves the control problem, while the outer loop updates world-modeling parameters for downstream objectives.
  • Scalability: The symplectic LQR solver replaces sequential matrix solves with parallel matrix operations, addressing accelerator efficiency constraints.The implementation is designed for high computational parallelism and throughput.
  • Approach: TTC-Net embeds finite-horizon LQR planning over internal representations and decodes optimal control actions as next-token representations.A TTC layer maps context-encoded latent states to optimal actions through linear dynamics and quadratic costs.
  • Results: TTC-Net consistently outperforms memory-only models on mathematical and symbolic reasoning, including up to +27.8% on MATH-500 and 2-3× Pass@8 gains on AMC and AIME.The architecture can be inserted as an adapter into pretrained LLMs without modifying the base architecture.

2 Background: Memory-Based Architectures

The background presents attention, linear RNNs, and related test-time-training systems as associative memory mechanisms that learn or retrieve information from token histories. These architectures differ in parameterization and update rules but share an online prediction objective over contextual representations.

  • Associative memory: Autoregressive models predict the next token by retrieving information associated with the preceding context.The background describes this as a common formulation across several sequence-processing architectures.
  • Memory units: A memory unit learns an online predictor from the observed prefix and returns its output for a query derived from the current token.Its forward pass encodes patterns in the prefix through an online self-supervised regression objective.
  • Memory units: The memory objective uses projected key and value views, reweights past tokens, regularizes the predictor, and evaluates it on a projected query.The resulting predictor output is produced at the current sequence position.
  • Architectural variants: Attention appears as a non-parametric regressor using uniform weights and no regularization, whereas linear RNNs use parametric linear mappings stored in fixed-size states.Attention traverses past-token caches, while linear RNN states track the history within a fixed-size memory block.
  • Architectural variants: Different linear RNN and SSM variants emerge by changing the state-update rule, token weighting, regularization, and optimization procedure.The background also notes that nonlinear predictors, alternative losses, decay, regularization, and advanced optimizers extend the memory-based family.

3 A Planning-Based Neural Architecture

TTC models next-token prediction as finite-horizon optimal control over latent states, using a tractable linear-quadratic formulation to choose the first action. Its differentiable and hardware-efficient solvers embed value-guided planning into neural architectures while reducing the sequential bottlenecks of classical LQR computation.

  • TTC formulation: TTC models next-token prediction as the first-step optimal decision in a finite-horizon Markov decision process over latent states.The initial latent state encodes prior context, while future latent states and token predictions are planned over horizon T.
  • TTC formulation: A linear dynamical system and quadratic reward provide a tractable approximation for modeling state transitions, process rewards, and terminal outcomes.Matrices A_t and B_t govern state evolution, while Q_t and R_t define quadratic costs over states and actions.
  • Learning and architecture: TTC layers decode a memory state into an optimal action and can be trained end-to-end by differentiating through a KKT formulation.The backward pass solves a structurally related dual LQR alongside the forward primal LQR to obtain parameter gradients.
  • Planning mechanism: Unlike memory-based layers, TTC optimizes future trajectories and computes a value function online at inference time.This gives each sequential modeling block a nested model-based reinforcement-learning objective.
  • Hardware co-design: Classical Riccati iteration is difficult to scale because sequential matrix inversions and memory traffic limit parallelism across long horizons.The paper identifies this compute- and I/O-bound behavior as a central challenge for training and inference over long sequences.
  • Hardware co-design: The symplectic solver replaces sequential Riccati recursion with cumulative matrix products and independently parallelizable inversions, requiring only one final dense inversion for the first action.This structure enables accelerator-friendly matrix multiplication and supports both primal and dual LQR solves.

4 TTC-Net: A Hybrid Model with TTC Layer

TTC-Net conditions LQR parameters on context and time, then inserts TTC layers into a hybrid memory architecture to support adaptable planning horizons and efficient inference.

  • Contextualization: TTC parameters are conditioned on the initial state and time step to capture context-specific dynamics, costs, and temporal variation.The parameterization uses context-informed coefficients and time modulation rather than a single fixed parameter set.
  • Hybrid with Memory Layers: TTC-Net inserts a TTC layer between attention and MLP modules, interleaving it with memory-based components every eight transformer blocks.Each TTC layer operates independently on tokens and receives context-rich attention features as its initial state.
  • Multi-Head Structure: The TTC layer uses multi-head processing, partitions input states into blocks, and combines head outputs after sharing basis matrices and selected scales.This structure is designed to reduce memory use while retaining distinct TTC parameters across heads.
  • Parameterization: Log-domain time modulation, direct parameterization of R_t^-1, and on-the-fly coefficient computation support numerical stability and efficient LQR evaluation.The implementation avoids repeatedly computing matrix inverses and does not instantiate time-dependent parameters in HBM.
  • Training Strategy: Mixed-horizon training samples bounded planning horizons to reduce distribution shift when the test-time horizon differs from the training horizon.The experiments use a mean horizon of 8, log-scale standard deviation 0.1, and maximum horizon 32.
  • Test-Time Scaling: Increasing the test-time planning horizon consumes additional FLOPs but enables deeper trajectory exploration and often more accurate solutions.The model generalizes from a maximum training horizon of 32 to T = 64 at test time, where accuracy continues to improve.

5 Experiments

Experiments evaluate TTC-Net on Sudoku and mathematical reasoning, showing gains over sequential-model and fine-tuned hybrid baselines. Ablations indicate that time heterogeneity, horizon sampling, and layer interleaving affect accuracy, generalization, and computational cost.

  • Sudoku Solving: TTC-Net surpasses the strongest runner-up, Transformer, by 2.8% in board-level accuracy on single-step Sudoku completion.It also demonstrates superior accuracy on multi-step Sudoku solving.
  • Math Reasoning: TTC-Net consistently outperforms other fine-tuned hybrid architectures across Math-500, AMC, AIME 2024, and AIME 2025.The base model has zero accuracy on both AIME datasets, whereas TTC-Net shows performance emergence; it also achieves large Pass@8 gains.
  • Test-Time Scaling: Reasoning accuracy consistently improves as the TTC planning horizon increases, including generalization from a maximum training horizon of 32 to T=64 at test time.Longer horizons increase test-time computation while enabling deeper trajectory exploration and emphasis on long-term objectives.
  • Ablation Studies: Time-homogeneous TTC underperforms time-heterogeneous TTC and degrades further when the planning horizon changes at test time.The results support time-heterogeneous dynamics and costs for modeling complex latent-space dynamics and generalizing across horizons.
  • Ablation Studies: Fixed-horizon training performs comparably at the matched horizon but fails to generalize or improve at larger test-time horizons.Uniform horizon sampling matches the full model’s accuracy but nearly doubles the average training horizon and substantially increases training cost relative to Poisson log-normal sampling.
  • Ablation Studies: Increasing TTC layers can improve accuracy monotonically, but evenly interleaving fewer TTC layers with attention is more cost-effective than stacking more layers.The comparison between 8:1 and 16:2 supports distributing TTC layers throughout the network.

6 Conclusion

TTC-Net embeds optimal-control planning into memory-based language models, using TTC layers to plan latent-state trajectories before next-token decoding. The paper also develops scalable LQR solvers, while identifying unresolved questions about multilayer dynamics and broader evaluation.

  • Conclusion: TTC-Net formulates reasoning as optimal control over internal representations, enabling future-trajectory planning before decoding the next token.The architecture augments memory-based language models with TTC layers.
  • Conclusion: TTC-Net reframes test-time learning as structured decision making that unifies memory, world modeling, reinforcement-learning objectives, and long-horizon planning.This broader framing places test-time control within a single architecture rather than treating planning as external to the model.
  • Limitations and Future Works: The joint interaction and dynamics represented by multiple TTC layers remain theoretically unclear despite the interpretable optimization objective of a single TTC layer.The paper identifies rigorous multilayer analysis as an open direction.
  • Limitations and Future Works: Future work should examine more expressive latent dynamics and reward models and evaluate TTC across larger-scale models and all training stages.These extensions must remain compatible with hardware-level optimization constraints.
  • Conclusion: A symplectic LQR solver provides an alternative finite-horizon solution method with the same O(Td^3) time and O(d^2) memory complexity as Riccati iteration.The symplectic algorithm uses one reverse and one forward iteration to solve the state, co-state, and control trajectories.

B Hardware Co-Designs and Optimization for TTC

The hardware–algorithm co-design reformulates LQR computation for parallel execution and fused CUDA implementation. The section also situates TTC among world-model, continuous-CoT, and test-time learning approaches while noting numerical-stability constraints.

  • Hardware Co-Designs and Optimization for TTC: Fused CUDA kernels stream LQR parameters through on-chip SRAM and reduce HBM access by combining tensor operations.The implementation also fuses primal and dual forward iterations for gradient computation.
  • Hardware Co-Designs and Optimization for TTC: The solver normalizes cumulative symplectic products row-wise to control numerical growth without changing the final computed result.This addresses overflow and instability caused by reciprocal-pair eigenvalues.
  • Comparison with Test-Time Discovery: Compared with Test-Time Discovery, TTC uses model-based value-based planning over hidden states without modifying model weights.The paper describes the two approaches as complementary.
  • Connection to World Models: TTC generates context-conditioned dynamics and costs, forming a simplified world model trained through supervision on predicted actions.The model emphasizes task-relevant dynamics and objectives rather than high-fidelity environment simulation.

D.1 Experiment Details

The experiments evaluate TTC across reasoning settings using controlled model, optimization, planning-horizon, and data configurations. Math evaluation samples eight solutions for AMC and AIME but uses greedy decoding for Math-500.

  • Sudoku Solving: The Sudoku experiments use a 32-layer model with 4 attention heads, embedding dimension 128, and 4 parallel TTC heads.The TTC hidden and control dimension is d=16, with r=16 basis matrices.
  • Math Reasoning: Math models use 16 TTC heads with state-space dimension d=16 and stochastic planning horizons averaging 8 steps over a 1–32-step support.The math setup uses AdamW with global batch size 96 and learning rate 2 × 10^-5.
  • Math Reasoning: AMC and AIME evaluation samples 8 solutions per problem at temperature 0.6, whereas Math-500 uses greedy decoding for one solution per problem.
  • Reasoning Data Curation: The curated fine-tuning data contains roughly 4 million objectively checkable questions spanning mathematics, academic subjects, science, medicine, finance, and coding.Recent datasets were selected partly to mitigate potential leakage and contamination.
  • Efficiency Benchmark: Efficiency benchmarks measure complete forward-and-backward throughput and memory footprint on an NVIDIA H200 across planning horizons and batch sizes.The reported throughput uses median runtime with 20–80% percentile statistics.

D.2 Comparison with Decoding-Based Scaling

TTC-Net is compared with decoding-based test-time scaling and explicit chain-of-thought prompting. The reported comparisons emphasize higher accuracy with fewer generated tokens and complementary gains from TTC and explicit CoT.

  • Comparison with Decoding-Based Scaling: TTC-Net achieves higher accuracy than decoding-based methods while generating substantially fewer tokens.External search methods require more than 20K generated tokens on average in the reported comparison.
  • TTC and Explicit CoT: TTC improves AMC accuracy both without explicit CoT and with explicit CoT.Without explicit CoT, accuracy rises from 3.92% to 7.83%; with CoT, it rises from 20.78% to 23.34%.
  • Interpretation: TTC is presented as an additional latent planning mechanism rather than a replacement for chain-of-thought reasoning.The paper states that explicit CoT remains important for difficult problems.

D.3 Evaluation beyond Reasoning Benchmarks

Beyond mathematics, TTC-Net improves several reasoning-heavy and code-generation evaluations while remaining competitive on general language-understanding benchmarks.

  • The evaluation covers commonsense reasoning, reading comprehension, knowledge-intensive multiple-choice evaluation, and code generation.
  • TTC-Net gains on ARC-Challenge, MMLU, and HumanEval, which assess long-horizon reasoning or program synthesis.
  • TTC-Net remains competitive with the raw SFT baseline on HellaSwag and WinoGrande.These results suggest that the planning module preserves general language-modeling ability learned in the base models.

D.4 Evaluation on Qwen Models

TTC-Net transfers to the Qwen2.5-Math-7B backbone and improves over standard full-parameter fine-tuning across all evaluated mathematics benchmarks.

  • TTC-Net improves over standard full-parameter fine-tuning across all evaluated math benchmarks on Qwen2.5-Math-7B.
  • TTC layers are modular and can be inserted into transformer-based architectures without relying on a specific backbone.

D.5 Computational Analysis

TTC-Net uses hardware-level optimization to limit the inference cost of finite-horizon planning. Throughput declines as the planning horizon grows, but the model remains practical at T=128.

  • TTC-Net reduces memory I/O traffic with a hardware-friendly LQR solver and fused CUDA kernels.These components address the overhead that naive finite-horizon optimal-control implementations would introduce, especially for large planning horizons.
  • At T=8, throughput decreases from 47.16 to 45.77 tokens/s compared with a standard transformer.
  • Increasing the planning horizon gradually reduces throughput, while the model remains practical even at T=128.The reported behavior reflects a trade-off in which larger horizons allocate more computation to latent planning.

D.6 Analysis of Latent Trajectories

Latent-trajectory analyses indicate that TTC states become more informative over planning steps and can be aligned with explicit chain-of-thought representations through monotone-path matching.

  • Linear Probing on Sudoku: Sudoku cell-level decoding accuracy rises from 42.50% at step 0 to 50.38% at step 8.The TTC states receive no cell-level supervision during training; only the final outcome is used for prediction.
  • Linear Probing on Sudoku: The increasing decoding accuracy indicates that TTC planning progressively refines latent representations toward the terminal solution.
  • Latent Trajectory Alignment with CoT: A monotone-path alignment score matches each TTC latent state to a consecutive chunk of CoT token representations.The maximization over monotone paths can be solved using dynamic programming, allowing comparison of same-question and different-question pairs.
  • Latent Trajectory Alignment with CoT: The alignment analysis compares TTC latent trajectories with CoT representations derived from the same question and from different question pairs on AMC.
Loading 2603.09221v2…