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

cs.LGMar 10, 2026

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

Authors: Peihao Wang, Shan Yang, Xijun Wang, Tesi Xiao, Xin Liu, Changlong Yu, Yu Lou, Pan Li, +3 more

Organizations: Amazon · The University of Texas at Austin · Georgia Institute of Technology · University of Maryland, College Park · University of Pennsylvania

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

Figures & tables

Figure 1: Memory-only prediction vs. unified memory-control planning. Memory-based models resemble human System 1 processing, relying on associative retrieval and producing fixed responses. Inspired by System 2 cognition, we model reasoning as an explicit optimal control problem and propose TTC-Net , a unified architecture that integrates test-time control ( TTC ) layers to encode value functions during sequential modeling, enabling planning before prediction.

Figure 2: Test-Time Control (TTC) layer. To predict the next token, TTC layers think through a receding-horizon LQR problem with linear state evolution and a cost function. A hardware-efficient solver computes optimal actions using structured matrix operations, enabling test-time control with minimal inference overhead.

Figure 3: Benchmarking running speed and memory of different LQR solvers . Throughput is reported in TFLOPs/s on a logarithmic scale. Among all evaluated solvers, our method achieves over a 10x higher throughput while maintaining constant memory cost w.r.t. horizon. Zero throughput indicates an out-of-memory error during execution.

Figure 4: Overview of TTC-Net. We construct a hybrid model by inserting a TTC layer between attention and MLP.

Methods	Single-Step		Multi-Step
Methods	Board Acc.	Cell Acc.	Board Acc.	Cell Acc.
Transformer	58.50	86.54	90.10	94.08
Mamba	54.60	85.50	88.60	91.29
Mamba2	55.50	85.10	87.20	90.52
GDN	57.30	87.19	89.80	93.70
Samba	57.20	87.99	90.40	94.61

Table 1: Accuracy (%) on Sudoku reasoning. Best method (ours) is marked in bold while the runner-up is underlined .

Models	MATH-500	AMC		AIME24		AIME25
Models	MATH-500	Acc@8	Pass@8	Acc@8	Pass@8	Acc@8	Pass@8
Base model	25.00	6.63	31.32	0.00	0.00	0.00	0.00
Full Finetuning	46.80	20.78	46.98	1.67	6.67	0.00	0.00
+ MLP ( Houlsby et al., 2019 )	40.80	18.67	33.73	0.00	0.00	0.00	0.00
+ Attention ( Vaswani et al., 2017 )	47.00	20.48	44.58	0.42	3.33	1.25	6.67
+ RetNet ( Sun et al., 2023 )	42.60	19.58	39.76	2.50	13.33	0.00	0.00

Table 2: Accuracy (%) on math reasoning datasets compared with baseline methods, including test-time memorization approaches. We mark the best performer in bold and the runner-up with underline . Our TTC-Net achieves the highest level of accuracy consistently.

Figure 5: Test-time scaling with TTC layers. TTC allows scaling test-time compute to improve performance by enlarging the planning horizon $T$ .

Variants	MATH-500
Base model	25.00
Parameterization
Homo. ( $T_{test}=8$ )	48.40
Homo. ( $T_{test}=16$ )	45.70
Horizon sampling
Fixed ( $T_{test}=8$ )	50.60

Table 3: Ablation studies. Accuracy (%) on MATH-500 dataset.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

	Fast weight	Slow weight
SSL	DeltaNet ( Yang et al., 2024b ) , GDN ( Yang et al., 2024a ) Titans ( Behrouz et al., 2024 ) , Nested Learning ( Behrouz et al., 2025b ) MesaNet ( von Oswald et al., 2025 ) , GatedKalmanNet ( Peng et al., 2025b )	TTT-MLP ( Sun et al., 2024 ) E2E-TTT ( Tandon et al., 2025 )
RL	TTC (Ours)	TTD ( Yuksekgonul et al., 2026 )

Appendix

Table 4: Taxonomy of various test-time learning methods.

Method	Acc.	Avg. Tokens
BoN	22.43	20731.10
ToT	21.08	22412.02
TTC-Net	23.34	2801.50

Appendix

Table 5: Comparison with decoding-based test-time scaling methods.s TTC-Net improves accuracy while requiring substantially fewer generated tokens.

Model	w/ CoT	w/o CoT
w/o TTC	20.78	3.92
w/ TTC	23.34	7.83

Appendix

Table 6: Ablation study between TTC and explicit CoT prompting on AMC.

Model	HellaSwag	WinoGrande	ARC-C	MMLU	HumanEval
LLaMA + SFT	54.01	74.74	51.79	59.31	32.32
TTC-Net	53.75	72.09	59.22	64.02	49.39

Appendix

Table 7: Evaluation beyond math reasoning benchmarks. TTC-Net improves several reasoning-heavy and code-related tasks while maintaining competitive performance on general language understanding benchmarks.

Model	MATH-500	AMC	AIME24	AIME25
Qwen2.5-Math-7B	43.8	33.0	6.7	6.7
+ Fine-tuning	74.4	51.7	20.4	18.7
+ TTC-Net	75.6	55.4	22.9	20.8

Appendix

Table 8: Results on Qwen2.5-Math-7B. TTC-Net yields consistent improvements over standard fine-tuning on other LLM backbones.

Model	Tokens/s
Transformer	47.16
TTC-Net ( $T=8$ )	45.77
TTC-Net ( $T=16$ )	44.61
TTC-Net ( $T=64$ )	41.43
TTC-Net ( $T=128$ )	36.18

Appendix

Table 9: Inference throughput of Transformer and TTC-Net on an NVIDIA H100 GPU.

Step	Cell Acc.
0	42.50
2	42.96
4	46.20
8	50.38

Appendix

Table 10: Linear probing accuracy of TTC latent states on Sudoku.

Pairing	Similarity
Same question	0.534
Random pairs	0.212

Appendix

Table 11: Trajectory-level alignment between TTC latent dynamics and explicit CoT representations.

Explore similar work

State of Thought Enables Endogenous Reasoning
Sep 13, 2026Zhiren Gong, Yikun Hou, Zihao Zeng +3Efficient Language Model Reasoning Implicit Reasoning in Language Models
College of Computing and Data Science, Nanyang Technological University, Singapore · Interdisciplinary Graduate Programme, Nanyang Technological University, Singapore · Information Science & Engineering, KTH Royal Institute of Technology, Sweden +1
Parason: Revealing Subtask and Trial Parallelism in LLM Reasoning
Aug 25, 2026Zhengyang Zhang, Zijian Zhang, Jiaxuan Gao +4Efficient Language Model Reasoning LLM Reasoning
Tsinghua University · NVIDIA · Independent Researcher +1
$\nabla$ -Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space
Mar 5, 2026Peihao Wang, Ruisi Cai, Zhen Wang +4LLM Reasoning Test-Time Optimization
The University of Texas at Austin · UC San Diego · TTIC +1

Sep 13, 2026J/K move · Enter open · S save

State of Thought Enables Endogenous Reasoning

Zhiren Gong, Yikun Hou, Zihao Zeng, Ming Xiao+2

College of Computing and Data Science, Nanyang Technological University, Singapore · Interdisciplinary Graduate Programme, Nanyang Technological University, Singapore · Information Science & Engineering, KTH Royal Institute of Technology, Sweden+1

Efficient Language Model Reasoning Implicit Reasoning in Language Models LLM Inference Acceleration

Open paper AskPDF arXiv

Test-time compute has emerged as a major approach to improving the capabilities of Large Language Models (LLMs). However, existing test-time reasoning paradigms rely heavily on externally imposed control, either through fixed reasoning programs or through costly expansion in constrained search spaces, limiting both generalization and efficiency. We propose State of Thought (SoT), a new reasoning paradigm that enables endogenous reasoning in LLMs, with the model's internal reasoning state governing how reasoning unfolds. Concretely, SoT extracts a compact dynamics-geometric state from the model's internal information transfer and uses a 582-parameter controller on frozen backbones to selectively activate historical reasoning support useful under the current reasoning state, framing reasoning as a state-conditioned process over evidence rather than an externally prescribed token chain. Across quantitative (1.34x), general (1.62x), symbolic-and-code (1.76x), and long-context (2.51x) reasoning on 3 LLMs and 16 datasets, SoT consistently improves mean-baseline accuracy while reducing generated tokens by 62.6% and end-to-end latency by 44.6%. Across 2 VLM scales and 3 reasoning tasks, it improves mean accuracy by 3.8 points over reasoning baselines, with 74.9% fewer completion tokens and 73.5% lower latency than search-based methods. Under constrained access, SoT retains 38.2%/36.5% mean accuracy gains in training-free/embedding-only settings, while trajectory-only judging reaches 84.1% agreement across 3 API models. Together, endogenous state-driven reasoning provides a generalizable and efficient alternative.

Abstract

Explore similar work

State of Thought Enables Endogenous Reasoning

Parason: Revealing Subtask and Trial Parallelism in LLM Reasoning

∇\nabla∇-Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space

$\nabla$ -Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space