Research map / Large language models / Reasoning and chain of thought

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

ICML 2026 · Large language models › Reasoning and chain of thought · approach: Efficient and adaptive thinking · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Test-Time Control (TTC) layers that run a differentiable finite-horizon LQR planner over latent states inside LLM blocks, with a symplectic-iteration fused CUDA solver, and adds them as adapters to Llama-3-Instruct-7B for math reasoning gains.

Problem

Memory-based sequence models (attention, linear RNNs, SSMs) predict next tokens via associative recall and lack an architectural mechanism for goal-directed planning; RL and test-time training keep planning external to the model.

Approach

Embed a finite-horizon LQR planning problem as a neural layer whose first-step optimal action is the next-token representation, with a value function internal to the layer, trained end-to-end via a KKT-based dual LQR and solved with a parallel symplectic iteration instead of sequential Riccati recursion.

Key results

• TTC-Net reaches 52.80% on MATH-500 vs 25.00% base and 47.80% for the best hybrid baseline (GDN), with the reported gain up to +27.8%.
• On AIME24 Pass@8 TTC-Net scores 20.00 vs 0.00 for base model and 13.33 for RetNet; AIME25 Pass@8 20.00 vs 6.67 best baseline.
• The symplectic iteration solver achieves over 10x higher throughput than Riccati- and KKT-based solvers while keeping memory constant in horizon.

Surprising findings

• Hybrid memory adapters (MLP, RetNet, Mamba, GDN, MesaNet) fail to consistently beat plain SFT on the base model, while the TTC layer does.
• Base model scores 0% on AIME24 and AIME25 but TTC-Net shows nonzero Pass@8, which the authors read as extending the reasoning boundary of the backbone.

Limitations

• Theoretical understanding of how multiple TTC layers interact in a deep transformer is unclear.
• Evaluation limited to a 7B model and small-scale training; larger-scale evaluation left open.
• Linear dynamics and quadratic costs may limit expressivity of latent dynamics and reward modeling.

Experimental setup

Datasets: OpenThoughts2-114K Sudoku 10k boards (hard dataset)

Benchmarks: MATH-500 AMC AIME24 AIME25 Sudoku

Metrics: accuracy Acc@8 Pass@8 board accuracy cell accuracy throughput TFLOPs/s

Baselines: Transformer Mamba Mamba2 GDN Samba RetNet MesaNet Full Finetuning Attention adapter

Builds on: Riccati recursion DeltaNet Mamba Gated DeltaNet Samba Test-time training

Closest papers by meaning

• MILES: Modular Instruction Memory with Learnable Selection for Self-Improving LLM Reasoning (NeurIPS 2026 · similarity 0.90)
• Reason for Future, Act for Now: A Principled Architecture for Autonomous LLM Agents (ICML 2024 · similarity 0.89)
• TTRL: Test-Time Reinforcement Learning (NeurIPS 2025 · similarity 0.89)
• Optimizing Test-Time Compute via Meta Reinforcement Finetuning (ICML 2025 · similarity 0.89)
• LLMs Improving LLMs: Agentic Discovery for Test-Time Scaling (NeurIPS 2026 · similarity 0.89)
• Satori: Reinforcement Learning with Chain-of-Action-Thought Enhances LLM Reasoning via Autoregressive Search (ICML 2025 · similarity 0.89)