Pith. sign in

REVIEW 2 major objections 5 minor 53 references

Reasoning is formulated as finite-horizon optimal control inside the model, so language models plan before they predict.

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

T0 review · grok-4.5

2026-07-15 12:07 UTC pith:ICQ5CP6O

load-bearing objection Solid hardware-efficient LQR adapter with real gains on math; the 'planning' story is plausible but not yet isolated from capacity. the 2 major comments →

arxiv 2603.09221 v2 pith:ICQ5CP6O submitted 2026-03-10 cs.LG

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
show 3 more authors
Zhangyang Wang Ming Lin Ren\'e Vidal
This is my paper
classification cs.LG
keywords test-time controllinear-quadratic regulatorstate-space modelsLLM reasoningvalue functionsymplectic iterationhybrid architectureplanning before prediction
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

Modern language models mostly work by retrieving and matching stored context, which is good for fluent next-token prediction but weak for deliberate multi-step problem solving. This paper argues that the missing piece is not more memory or more external training, but an architectural layer that plans. It inserts a Test-Time Control (TTC) layer that treats the next-token decision as the first action of a short-horizon Linear-Quadratic Regulator (LQR) problem solved over latent states at inference time. Dynamics and costs are synthesized from the current context, a value function is computed on the fly, and the optimal first action becomes the next representation. A hardware-efficient symplectic solver, fused into CUDA kernels, keeps the cost low enough to run as an adapter inside pretrained models. On Sudoku and contest math, the resulting hybrid networks outperform pure memory adapters and full fine-tuning, with large gains on MATH-500, AMC, and AIME. The claim is that embedding optimal control as a native architectural component lets models reason over future trajectories without treating planning as an external procedure.

Core claim

The paper establishes that reasoning can be internalized as finite-horizon LQR planning over latent states: a TTC layer synthesizes context-dependent linear dynamics and quadratic costs, solves for the optimal first action at test time, and thereby encodes a value function that lets the network plan before it predicts. When these layers are interleaved with attention as adapters in pretrained LLMs, mathematical reasoning accuracy rises substantially relative to memory-only baselines.

What carries the argument

The Test-Time Control (TTC) layer: a differentiable finite-horizon LQR solver that takes a context-encoded initial state, produces time-heterogeneous dynamics and costs, computes the optimal first-step control via symplectic iteration, and returns that action as the next-token representation, with gradients obtained by solving a related dual LQR.

Load-bearing premise

The claim rests on the premise that a short-horizon linear-quadratic model whose parameters are read from a single context vector is rich enough to stand in for the multi-step reasoning needed by contest math and Sudoku.

What would settle it

If, on the same fine-tuning recipe and parameter budget, replacing the TTC layer with a capacity-matched memory adapter (or a non-planning MLP) erases the reported MATH-500 and Pass@8 gains, or if lengthening the planning horizon past the training distribution systematically degrades rather than improves accuracy, the architectural planning benefit is falsified.

Watch this falsifier — get emailed when new claim-graph text bears on it.

If this is right

  • Planning becomes a native forward-pass operation rather than an external RL or search loop, so longer test-time horizons can be dialed up without changing model weights.
  • Hybrid memory-plus-control blocks can be dropped into existing transformers as zero-init adapters and still unlock measurable contest-math gains.
  • Hardware-efficient symplectic LQR solvers make value-function layers practical at the scale of modern LLMs.
  • Test-time compute can be spent on deeper latent planning instead of only longer token-level search.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • If short-horizon LQR is already useful, richer but still tractable controllers (time-varying or mildly nonlinear) may further close the gap on harder multi-step domains.
  • The same adapter pattern could be tried on non-math sequential tasks that reward foresight, such as long-horizon code or multi-turn tool use, to test whether the benefit is domain-specific.
  • Latent trajectories that align with explicit chain-of-thought suggest the control layer may be compressing token-level reasoning into a compact state-space plan.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

2 major / 5 minor

Summary. The paper argues that modern LLMs are limited by associative-memory architectures (System-1 style) and that reasoning can be internalized as finite-horizon optimal control. It introduces the Test-Time Control (TTC) layer, which, given a context-encoded latent state h0, synthesizes time-heterogeneous linear dynamics and quadratic costs, solves the resulting LQR via a symplectic-iteration algorithm (Theorem 3.3 / Algorithm 3) implemented as a fused CUDA kernel, and emits the first-step optimal action u*_1 as the next-token representation. Differentiability is obtained by a dual LQR derived from the KKT system (Theorem 3.2). TTC layers are inserted as adapters (every 8 blocks) into pretrained transformers, yielding hybrid TTC-Nets. Empirically, the method improves Sudoku board/cell accuracy over Transformer/Mamba/GDN/Samba baselines and, after SFT on OpenThoughts2 + curated data, raises MATH-500 accuracy by up to 27.8 points over the Llama-3-7B base (and outperforms matched memory adapters) while producing 2–3 imes Pass@8 gains on AMC/AIME; test-time horizon scaling further improves accuracy.

Significance. If the central claim holds—that embedding a value-function / LQR planner as a native architectural component, rather than external RL or pure test-time memorization, is what drives the observed reasoning gains—the work supplies a concrete, hardware-efficient primitive for System-2 computation inside LLMs. The symplectic reformulation, KKT gradient derivation, structured diagonal/low-rank parameterization, and fused CUDA kernels are carefully executed and constitute a genuine systems contribution that makes the idea practical. Consistent gains across Sudoku, MATH-500, AMC, AIME, a second backbone (Qwen2.5-Math), and non-math LM benchmarks, together with progressive latent probing and CoT-alignment analyses in the appendix, give the result real weight. The paper therefore has the potential to influence both architectural design and the broader conversation on how to internalize planning, world models, and nested RL objectives.

major comments (2)
  1. [§5.2, Table 2] The headline interpretation—that gains arise specifically from solving the LQR for a value-optimal first action u*_1 (Prop. 3.1 / Thm. 3.3) rather than from the mere insertion of a context-conditioned multi-step linear dynamical system of comparable depth and parameter count—is not isolated by the experiments. Table 2 and the adapter suite (MLP, Attention, RetNet, Mamba, GDN, MesaNet) compare against other memory mechanisms but never reuse the Sec. 4 contextualization that produces {(A_t,B_t,Q_t,R_t)} from h0 and then replace only the optimal solve by a non-optimal map of matched structure (open-loop zero/random-action rollout of the same dynamics, a fixed non-Riccati gain, or a dense readout from h0 with the same free-parameter budget). Without that control, the MATH-500 / AMC / AIME lifts remain consistent with ordinary residual capacity addition; the “planning before prediction / nest
  2. [§3.1, Eqs. (5)–(7); §4 contextualization] The modeling premise that a short-horizon (T_train ~ PLN mean 8, max 32; T_test = 8–16) linear-quadratic regulator whose dynamics and costs are synthesized from a single context vector is a sufficiently faithful model of multi-step contest-math and Sudoku reasoning (Eqs. 5–7 and the power-series contextualization in Sec. 4) remains a strong and largely untested assumption. While Fig. 5 shows beneficial test-time horizon scaling and App. D.6 linear probes / CoT-alignment scores indicate progressive refinement, these do not rule out that the benefit collapses to ordinary multi-step capacity once the true latent dynamics become strongly nonlinear or the relevant horizon far exceeds the trained range. A clearer statement of this limitation, or an experiment that deliberately degrades the optimality of the solve while preserving the dynamics, would strengthen the paper.
minor comments (5)
  1. [Abstract] Abstract and introduction report “+27.8% on MATH-500”; the absolute point gain from the 25.00 base is 27.8, but the wording can be read as a relative percentage. Clarify “points” versus “percent relative improvement.”
  2. [Fig. 3] Figure 3 throughput is plotted on a log scale and reports “over 10× higher throughput”; the caption and surrounding text should also state the absolute TFLOPs/s numbers and the precise hardware (H200) so that the comparison is reproducible without reading the appendix.
  3. [§4 Multi-Head Structure] The multi-head and basis-sharing design (r=16, d=16 per head) is described only briefly; a short paragraph or table listing the exact parameter count of a TTC adapter versus the memory adapters used in Table 2 would make the “comparable capacity” claim fully transparent.
  4. [§5.1, Alg. 6] In the Sudoku multi-step protocol (Alg. 6) the model is allowed to fill only the most confident cell each iteration; it would be useful to report whether single-step board accuracy already correlates strongly with multi-step success, or whether the iterative procedure is essential for the reported gains.
  5. [§3 and App. A] A few minor notational inconsistencies appear (e.g., h0 versus h_init, Q_0 = 0 stated in different places). A single consistent convention throughout Sec. 3 and Appendix A would improve readability.

Circularity Check

0 steps flagged

No significant circularity: LQR identities, symplectic solver, and KKT gradients are standard self-contained derivations; gains are measured on external held-out benchmarks.

full rationale

The paper's derivation chain (reasoning as finite-horizon LQR over latent states, Prop. 3.1 Riccati, Thm. 3.3 symplectic iteration, Thm. 3.2 dual-LQR gradients via KKT) consists of classical optimal-control identities that are re-derived with explicit proofs in Appendix A and do not reduce to the paper's own empirical inputs or fitted constants. Contextualization of {(A_t, B_t, Q_t, R_t)} from h_0 (Sec. 4) is an architectural design choice, not a uniqueness claim forced by self-citation. Performance numbers (+27.8% MATH-500, 2-3x Pass@8 on AMC/AIME, Sudoku board/cell accuracies) are obtained by supervised fine-tuning or training from scratch and evaluated on held-out external benchmarks against memory/TTT adapters; nothing is fitted on a subset and then reported as a forced prediction. Background citations to SSM/TTT/nested-learning literature supply context for the memory baseline but are not load-bearing for the LQR formulation or the reported lifts. The absence of an ablation that isolates the optimality of the solve from multi-step linear capacity is an experimental-design gap, not circularity. The work is therefore self-contained against external evidence.

Axiom & Free-Parameter Ledger

5 free parameters · 3 axioms · 2 invented entities

The central empirical claim rests on the modeling choice that latent reasoning is well-approximated by a short-horizon LQR whose parameters are generated from a single context vector, plus a handful of architectural free parameters (horizon distribution, basis rank, interleaving ratio) that are chosen by the authors rather than derived. The mathematical identities themselves are standard; the free parameters and the LQR-as-reasoning axiom are the load-bearing extras.

free parameters (5)
  • planning-horizon distribution (PLN mean T_mu=8, sigma=0.1, max 32)
    Chosen by hand; mixed-horizon training is required for test-time horizon generalization, yet the distribution itself is not derived from data or theory.
  • basis rank r=16 for Q^(i) and B^(i)
    Controls expressivity of the contextualized cost and dynamics; set empirically.
  • state/control dimension d=16 per head and number of heads
    Architectural free parameters that determine capacity of each TTC layer.
  • interleaving ratio Attn:TTC = 8:1
    Selected after ablation; affects both accuracy and compute.
  • test-time horizons T_test = 8/16/16 for MATH/AMC/AIME
    Chosen after training; performance continues to rise with longer horizons, so reported numbers depend on this choice.
axioms (3)
  • domain assumption A finite-horizon linear-quadratic regulator with context-dependent A_t, B_t, Q_t, R_t is a sufficiently expressive model of the latent MDP that underlies multi-step mathematical and constraint-satisfaction reasoning.
    Stated in Sec. 3.1 and used to justify replacing a general Bellman equation by Eqs. 5–7; no independent verification that the true reasoning dynamics are linear-quadratic.
  • ad hoc to paper Diagonal A_t and R_t together with low-rank bases for B and Q retain enough expressivity for the tasks considered.
    Adopted for hardware efficiency (Sec. 3.3); justified by analogy to diagonal SSMs but not proven for the control setting.
  • standard math Standard KKT and symplectic identities for finite-horizon LQR hold under the chosen positive-semidefinite and invertibility constraints.
    Classic optimal-control theory (Kalman, Bellman, Laub); used throughout Sec. 3 and Appendix A.
invented entities (2)
  • Test-Time Control (TTC) layer no independent evidence
    purpose: Maps a context-encoded latent state to the first-step optimal action of a synthesized LQR, thereby inserting a value function and planning step into the residual stream.
    The core architectural invention; independent evidence is the empirical accuracy lift and the open-sourceable CUDA kernel, but the entity itself is defined only inside this paper.
  • Symplectic-iteration LQR solver fused as CUDA kernel no independent evidence
    purpose: Replaces sequential Riccati recursion by parallel matrix products and a constant number of dense inversions so that TTC can run at LLM scale.
    New algorithmic packaging of known symplectic structure; independent evidence would be third-party re-implementation or formal verification, neither of which is supplied.

pith-pipeline@v1.1.0-grok45 · 42503 in / 3333 out tokens · 32070 ms · 2026-07-15T12:07:55.396964+00:00 · methodology

0 comments
read the original 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.

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

53 extracted references · 38 linked inside Pith

  1. [1]

    xlstm: Extended long short-term memory.arXiv preprint arXiv:2405.04517,

    Beck, M., Pöppel, K., Spanring, M., Auer, A., Prudnikova, O., Kopp, M., Klambauer, G., Brandstetter, J., and Hochreiter, S. xlstm: Extended long short-term memory.arXiv preprint arXiv:2405.04517,

    Pith/arXiv arXiv

  2. [2]

    Titans: Learning to memorize at test time.arXiv preprint arXiv:2501.00663,

    Behrouz, A., Zhong, P., and Mirrokni, V. Titans: Learning to memorize at test time.arXiv preprint arXiv:2501.00663,

    Pith/arXiv arXiv

  3. [3]

    Atlas: Learning to optimally memorize the context at test time.URL https://arxiv

    Behrouz, A., Li, Z., Kacham, P., Daliri, M., Deng, Y., Zhong, P., Razaviyayn, M., and Mirrokni, V. Atlas: Learning to optimally memorize the context at test time.URL https://arxiv. org/abs/2505.23735, 2025a. Behrouz, A., Razaviyayn, M., Zhong, P., and Mirrokni, V. Nested learning: The illusion of deep learning architectures. arXiv preprint arXiv:2512.2469...

    Pith/arXiv arXiv

  4. [4]

    The riccati equation(book).Berlin and New York, Springer-Verlag, 1991, 347,

    Bittanti, S., LAUB, A., and WILLEMS, J. The riccati equation(book).Berlin and New York, Springer-Verlag, 1991, 347,

    1991

  5. [5]

    Learning phrase representations using rnn encoder-decoder for statistical machine translation.arXiv preprint arXiv:1406.1078,

    Cho, K., Van Merriënboer, B., Gulcehre, C., Bahdanau, D., Bougares, F., Schwenk, H., and Bengio, Y. Learning phrase representations using rnn encoder-decoder for statistical machine translation.arXiv preprint arXiv:1406.1078,

    Pith/arXiv arXiv

  6. [6]

    and Gu, A

    Dao, T. and Gu, A. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality.arXiv preprint arXiv:2405.21060,

    Pith/arXiv arXiv

  7. [7]

    De Luca, A. B. and Fountoulakis, K. Simulation of graph algorithms with looped transformers.arXiv preprint arXiv:2402.01107,

    Pith/arXiv arXiv

  8. [8]

    Reinforcement pre-training.arXiv preprint arXiv:2506.08007,

    Dong, Q., Dong, L., Tang, Y., Ye, T., Sun, Y., Sui, Z., and Wei, F. Reinforcement pre-training.arXiv preprint arXiv:2506.08007,

    Pith/arXiv arXiv

  9. [9]

    S., Liu, S.-Y., Van Keirsbilck, M., Chen, M.-H., Suhara, Y., et al

    Dong, X., Fu, Y., Diao, S., Byeon, W., Chen, Z., Mahabaleshwarkar, A. S., Liu, S.-Y., Van Keirsbilck, M., Chen, M.-H., Suhara, Y., et al. Hymba: A hybrid-head architecture for small language models.arXiv preprint arXiv:2411.13676,

    Pith/arXiv arXiv

  10. [10]

    Mom: Linear sequence modeling with mixture-of-memories.arXiv preprint arXiv:2502.13685,

    Du, J., Sun, W., Lan, D., Hu, J., and Cheng, Y. Mom: Linear sequence modeling with mixture-of-memories.arXiv preprint arXiv:2502.13685,

    arXiv

  11. [11]

    R., Kailkhura, B., Bhatele, A., and Goldstein, T

    Geiping, J., McLeish, S., Jain, N., Kirchenbauer, J., Singh, S., Bartoldson, B. R., Kailkhura, B., Bhatele, A., and Goldstein, T. Scaling up test-time compute with latent reasoning: A recurrent depth approach.arXiv preprint arXiv:2502.05171,

    Pith/arXiv arXiv

  12. [12]

    and Dao, T

    Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces.arXiv preprint arXiv:2312.00752,

    Pith/arXiv arXiv

  13. [13]

    Efficiently modeling long sequences with structured state spaces.arXiv preprint arXiv:2111.00396, 2021a

    Gu, A., Goel, K., and Ré, C. Efficiently modeling long sequences with structured state spaces.arXiv preprint arXiv:2111.00396, 2021a. Gu, A., Johnson, I., Goel, K., Saab, K., Dao, T., Rudra, A., and Ré, C. Combining recurrent, convolutional, and continuous-time models with linear state space layers.Advances in neural information processing systems, 34:572...

    Pith/arXiv arXiv

  14. [14]

    Openthoughts: Data recipes for reasoning models.URL https://arxiv

    Guha, E., Marten, R., Keh, S., Raoof, N., Smyrnis, G., Bansal, H., Nezhurina, M., Mercat, J., Vu, T., Sprague, Z., et al. Openthoughts: Data recipes for reasoning models.URL https://arxiv. org/abs/2506.04178,

    Pith/arXiv arXiv

  15. [15]

    Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning.arXiv preprint arXiv:2501.12948,

    Guo, D., Yang, D., Zhang, H., Song, J., Zhang, R., Xu, R., Zhu, Q., Ma, S., Wang, P., Bi, X., et al. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning.arXiv preprint arXiv:2501.12948,

    Pith/arXiv arXiv

  16. [16]

    and Schmidhuber, J

    Ha, D. and Schmidhuber, J. World models.arXiv preprint arXiv:1803.10122, 2(3),

    Pith/arXiv arXiv

  17. [17]

    Dream to control: Learning behaviors by latent imagination.arXiv preprint arXiv:1912.01603,

    Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. Dream to control: Learning behaviors by latent imagination.arXiv preprint arXiv:1912.01603,

    Pith/arXiv arXiv 1912

  18. [18]

    Temporal difference learning for model predictive control.arXiv preprint arXiv:2203.04955,

    Hansen, N., Wang, X., and Su, H. Temporal difference learning for model predictive control.arXiv preprint arXiv:2203.04955,

    Pith/arXiv arXiv

  19. [19]

    Td-mpc2: Scalable, robust world models for continuous control.arXiv preprint arXiv:2310.16828,

    Hansen, N., Su, H., and Wang, X. Td-mpc2: Scalable, robust world models for continuous control.arXiv preprint arXiv:2310.16828,

    Pith/arXiv arXiv

  20. [20]

    J., Wang, Z., Wang, D

    Hao, S., Gu, Y., Ma, H., Hong, J. J., Wang, Z., Wang, D. Z., and Hu, Z. Reasoning with language model is planning with world model.arXiv preprint arXiv:2305.14992,

    Pith/arXiv arXiv

  21. [21]

    Training large language models to reason in a continuous latent space.arXiv preprint arXiv:2412.06769,

    Hao, S., Sukhbaatar, S., Su, D., Li, X., Hu, Z., Weston, J., and Tian, Y. Training large language models to reason in a continuous latent space.arXiv preprint arXiv:2412.06769,

    Pith/arXiv arXiv

  22. [22]

    N., Prabhumoye, S., Kautz, J., Patwary, M., Shoeybi, M., Catanzaro, B., and Choi, Y

    Hatamizadeh, A., Akter, S. N., Prabhumoye, S., Kautz, J., Patwary, M., Shoeybi, M., Catanzaro, B., and Choi, Y. Rlp: Reinforcement as a pretraining objective.arXiv preprint arXiv:2510.01265,

    arXiv

  23. [23]

    Measuring mathematical problem solving with the math dataset.arXiv preprint arXiv:2103.03874,

    Hendrycks, D., Burns, C., Kadavath, S., Arora, A., Basart, S., Tang, E., Song, D., and Steinhardt, J. Measuring mathematical problem solving with the math dataset.arXiv preprint arXiv:2103.03874,

    Pith/arXiv arXiv

  24. [24]

    A., et al

    Kobayashi, S., Schimpf, Y., Schlegel, M., Steger, A., Wolczyk, M., von Oswald, J., Scherre, N., Maile, K., Lajoie, G., Richards, B. A., et al. Emergent temporal abstractions in autoregressive models enable hierarchical reinforcement learning.arXiv preprint arXiv:2512.20605,

    arXiv

  25. [25]

    Longhorn: State space models are amortized online learners.arXiv preprint arXiv:2407.14207,

    Liu, B., Wang, R., Wu, L., Feng, Y., Stone, P., and Liu, Q. Longhorn: State space models are amortized online learners.arXiv preprint arXiv:2407.14207,

    Pith/arXiv arXiv

  26. [26]

    Large language model guided tree-of-thought.arXiv preprint arXiv:2305.08291,

    Long, J. Large language model guided tree-of-thought.arXiv preprint arXiv:2305.08291,

    Pith/arXiv arXiv

  27. [27]

    Rwkv-7" goose" with expressive dynamic state evolution.arXiv preprint arXiv:2503.14456, 2025a

    Peng, B., Zhang, R., Goldstein, D., Alcaide, E., Du, X., Hou, H., Lin, J., Liu, J., Lu, J., Merrill, W., et al. Rwkv-7" goose" with expressive dynamic state evolution.arXiv preprint arXiv:2503.14456, 2025a. Peng, L., Chattopadhyay, A., Zancato, L., Nunez, E., Xia, W., and Soatto, S. Gated kalmanet: A fading memory layer through test-time ridge regression....

    Pith/arXiv arXiv 2008

  28. [28]

    Samba: Simple hybrid state space models for efficient unlimited context language modeling.arXiv preprint arXiv:2406.07522,

    Ren, L., Liu, Y., Lu, Y., Shen, Y., Liang, C., and Chen, W. Samba: Simple hybrid state space models for efficient unlimited context language modeling.arXiv preprint arXiv:2406.07522,

    Pith/arXiv arXiv

  29. [29]

    Saunshi, N., Dikkala, N., Li, Z., Kumar, S., and Reddi, S. J. Reasoning with latent thoughts: On the power of looped transformers.arXiv preprint arXiv:2502.17416,

    Pith/arXiv arXiv

  30. [30]

    Deepseekmath: Pushing the limits of mathematical reasoning in open language models.arXiv preprint arXiv:2402.03300,

    Shao, Z., Wang, P., Zhu, Q., Xu, R., Song, J., Bi, X., Zhang, H., Zhang, M., Li, Y., Wu, Y., et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models.arXiv preprint arXiv:2402.03300,

    Pith/arXiv arXiv

  31. [31]

    The illusion of thinking: Understanding the strengths and limitations of reasoning models via the lens of problem complexity

    Shojaee, P., Mirzadeh, I., Alizadeh, K., Horton, M., Bengio, S., and Farajtabar, M. The illusion of thinking: Understanding the strengths and limitations of reasoning models via the lens of problem complexity. arxiv 2025.arXiv preprint arXiv:2506.06941, 10,

    Pith/arXiv arXiv 2025

  32. [32]

    Scaling llm test-time compute optimally can be more effective than scaling model parameters.arXiv preprint arXiv:2408.03314,

    Snell, C., Lee, J., Xu, K., and Kumar, A. Scaling llm test-time compute optimally can be more effective than scaling model parameters.arXiv preprint arXiv:2408.03314,

    Pith/arXiv arXiv

  33. [33]

    Retentive network: A successor to transformer for large language models.arXiv preprint arXiv:2307.08621,

    Sun, Y., Dong, L., Huang, S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, F. Retentive network: A successor to transformer for large language models.arXiv preprint arXiv:2307.08621,

    Pith/arXiv arXiv

  34. [34]

    Learning to (learn at test time): Rnns with expressive hidden states.arXiv preprint arXiv:2407.04620,

    Sun, Y., Li, X., Dalal, K., Xu, J., Vikram, A., Zhang, G., Dubois, Y., Chen, X., Wang, X., Koyejo, S., et al. Learning to (learn at test time): Rnns with expressive hidden states.arXiv preprint arXiv:2407.04620,

    Pith/arXiv arXiv

  35. [35]

    End-to-end test-time training for long context.arXiv preprint arXiv:2512.23675,

    Tandon, A., Dalal, K., Li, X., Koceja, D., Rød, M., Buchanan, S., Wang, X., Leskovec, J., Koyejo, S., Hashimoto, T., et al. End-to-end test-time training for long context.arXiv preprint arXiv:2512.23675,

    arXiv

  36. [36]

    Uncovering mesa-optimization algorithms in transformers.arXiv preprint arXiv:2309.05858,

    Von Oswald, J., Schlegel, M., Meulemans, A., Kobayashi, S., Niklasson, E., Zucchet, N., Scherrer, N., Miller, N., Sandler, M., Vladymyrov, M., et al. Uncovering mesa-optimization algorithms in transformers.arXiv preprint arXiv:2309.05858,

    Pith/arXiv arXiv

  37. [37]

    Mesanet: Sequence modeling by locally optimal test-time training.arXiv preprint arXiv:2506.05233,

    von Oswald, J., Scherrer, N., Kobayashi, S., Versari, L., Yang, S., Schlegel, M., Maile, K., Schimpf, Y., Sieberling, O., Meulemans, A., et al. Mesanet: Sequence modeling by locally optimal test-time training.arXiv preprint arXiv:2506.05233,

    Pith/arXiv arXiv

  38. [38]

    Wang, G., Li, J., Sun, Y., Chen, X., Liu, C., Wu, Y., Lu, M., Song, S., and Yadkori, Y. A. Hierarchical reasoning model.arXiv preprint arXiv:2506.21734, 2025a. Wang, K. A., Shi, J., and Fox, E. B. Test-time regression: a unifying framework for designing sequence models with associative memory.arXiv preprint arXiv:2501.12352, 2025b. Wang, P.-W., Donti, P.,...

    Pith/arXiv arXiv

  39. [39]

    Self-consistency improves chain of thought reasoning in language models.arXiv preprint arXiv:2203.11171,

    Wang, X., Wei, J., Schuurmans, D., Le, Q., Chi, E., Narang, S., Chowdhery, A., and Zhou, D. Self-consistency improves chain of thought reasoning in language models.arXiv preprint arXiv:2203.11171,

    Pith/arXiv arXiv

  40. [40]

    Looped transformers are better at learning learning algorithms.arXiv preprint arXiv:2311.12424, 2023a

    Yang, L., Lee, K., Nowak, R., and Papailiopoulos, D. Looped transformers are better at learning learning algorithms.arXiv preprint arXiv:2311.12424, 2023a. Yang, S., Wang, B., Shen, Y., Panda, R., and Kim, Y. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:2312.06635, 2023b. Yang, S., Kautz, J., and Hatamizadeh, ...

    Pith/arXiv arXiv

  41. [41]

    Does reinforcement learning really incentivize reasoning capacity in llms beyond the base model?URL https://arxiv

    17 Yue, Y., Chen, Z., Lu, R., Zhao, A., Wang, Z., Yue, Y., Song, S., and Huang, G. Does reinforcement learning really incentivize reasoning capacity in llms beyond the base model?URL https://arxiv. org/abs/2504.13837,

    Pith/arXiv arXiv

  42. [42]

    Learning to discover at test time.arXiv preprint arXiv:2601.16175,

    Yuksekgonul, M., Koceja, D., Li, X., Bianchi, F., McCaleb, J., Wang, X., Kautz, J., Choi, Y., Zou, J., Guestrin, C., et al. Learning to discover at test time.arXiv preprint arXiv:2601.16175,

    Pith/arXiv arXiv

  43. [43]

    Pretraining language models to ponder in continuous space.arXiv preprint arXiv:2505.20674,

    Zeng, B., Song, S., Huang, S., Wang, Y., Li, H., He, Z., Wang, X., Li, Z., and Lin, Z. Pretraining language models to ponder in continuous space.arXiv preprint arXiv:2505.20674,

    arXiv

  44. [44]

    T., and Tan, H

    Zhang, T., Bi, S., Hong, Y., Zhang, K., Luan, F., Yang, S., Sunkavalli, K., Freeman, W. T., and Tan, H. Test-time training done right.arXiv preprint arXiv:2505.23884,

    Pith/arXiv arXiv

  45. [45]

    Emergence of superposition: Unveiling the training dynamics of chain of continuous thought.arXiv preprint arXiv:2509.23365, 2025a

    Zhu, H., Hao, S., Hu, Z., Jiao, J., Russell, S., and Tian, Y. Emergence of superposition: Unveiling the training dynamics of chain of continuous thought.arXiv preprint arXiv:2509.23365, 2025a. Zhu, R.-J., Wang, Z., Hua, K., Zhang, T., Li, Z., Que, H., Wei, B., Wen, Z., Yin, F., Xing, H., et al. Scaling latent reasoning via looped language models.arXiv pre...

    arXiv

  46. [46]

    In the main text, we use𝒉 0 in place of𝒉 𝑖𝑛𝑖𝑡 without separate notations when no confusion arises

    18 A Theory Definition A.1(Linear-Quadratic Regulator (Kalman et al., 1960)).Given LQR system with parameters {(𝑨 𝑡,𝑩 𝑡,𝑸 𝑡,𝑹 𝑡,𝒓 𝑡 )}𝑇 𝑡=1 and initial state𝒉 0 =𝒉 𝑖𝑛𝑖𝑡 , the optimization problem can be formally written as: min 𝒉0,𝒉1,···,𝒉𝑇 𝒖1,···,𝒖𝑇 1 2 (𝒉⊤ 𝑡 𝑸𝑡 𝒉𝑡 +𝒖 ⊤ 𝑡 𝑹𝑡 𝒖𝑡 ) +𝒓 ⊤ 𝑡 𝒖𝑡 (13) s.t.𝒉 𝑡+1 =𝑨 𝑡+1 𝒉𝑡 +𝑩 𝑡+1 𝒖𝑡+1,𝒉 0 =𝒉 𝑖𝑛𝑖𝑡 (14) Note that w...

    1960

  47. [47]

    𝜕 𝜕𝜃 𝒙 𝜕 𝜕𝜃 𝝃 # =

    10:for𝑡=1,2, . . . , 𝑇do⊲forward symplectic iteration 11: 𝒉𝑡 𝝀𝑡 ← 𝑨𝑡 +𝑮 𝑡 (𝑨⊤ 𝑡 ) −1𝑸𝑡−1 −𝑮𝑡 (𝑨⊤ 𝑡 ) −1 −(𝑨 ⊤ 𝑡 ) −1𝑸𝑡−1 (𝑨⊤ 𝑡 ) −1 𝒉𝑡−1 𝝀𝑡−1 + −𝑩𝑡 𝑹 −1 𝑡 𝒓𝑡 0 12:𝒖 𝑡 ← −𝑹 −1 𝑡 (𝑩 ⊤ 𝑡 𝝀𝑡 +𝒓 𝑡 ) 13:end for 14:return{𝒖 𝑡 }𝑇 𝑡=1,{𝒉 𝑡 }𝑇 𝑡=1,{𝝀 𝑡 }𝑇 𝑡=0. Remark A.8.Solving LQR with Algorithm 3 only requires one reverse and one forward iteration to solve all {...

    2017

  48. [48]

    Note that𝑺 𝑡 follows a similar factorized form: 𝑺𝑡 = 𝑰−𝑩 𝑡 𝑹 −1 𝑡 𝑩⊤ 𝑡 0𝑰 𝑨𝑡 0 0(𝑨 −1 𝑡 )⊤ 𝑰0 −𝑸𝑡−1 𝑰 , allowing accumulating 𝑺𝑡 via streaming 𝑸𝑡−1 , 𝑨𝑡, 𝑹𝑡, and 𝑩𝑡 in order. For gradient computation, we fuse two forward symplectic iterations for both primal and dual LQRs together, to efficiently obtain and combine trajectories {(𝒉 𝑡,𝒖 𝑡,𝝀 𝑡 )} and{( e𝒉𝑡,...

    2017

  49. [49]

    Although [𝒀 1,𝒀 2] are partitioned row-wise and distributed across CUDA blocks, this normalization can be performed independently within each kernel without synchronization

    Based on this observation, we construct a diagonal matrix 𝑫 with entries 𝑫𝑖𝑖 =max{∥𝒀 1,𝑖 ∥1,∥𝒀 2,𝑖 ∥1} and apply 𝑫 −1 to [𝒀 1,𝒀 2] immediately after each step. Although [𝒀 1,𝒀 2] are partitioned row-wise and distributed across CUDA blocks, this normalization can be performed independently within each kernel without synchronization. Mixed Precision.The TTC...

    2018

  50. [50]

    To the best of our knowledge, it is the only other approach that explores online RL–based adaptation at test time

    is a recent work that came out during the preparation of this paper. To the best of our knowledge, it is the only other approach that explores online RL–based adaptation at test time. However, the technical approaches differ substantially: TTD performs model-free, policy-based RL by updating the model parameters (i.e., slow weights), whereas our method ad...

    2024

  51. [51]

    world”, upon which it performs planning and decision-making. All parameters of this “world

    Table 4: Taxonomy of various test-time learning methods. Connection to World Models.TTC generates the environment parameters {(𝑨 𝑡,𝑩 𝑡,𝑸 𝑡,𝑹 𝑡 )}𝑇 𝑡=1 conditioned on the input context, and then solves a control problem defined over these parameters to produce the output. The resulting linear state transitions and quadratic cost together constitute a simpl...

    2018

  52. [52]

    thinking

    has been proposed to enable LLMs to reason beyond discrete tokens by representing intermediate reasoning traces as latent vectors. The general framework of continuous CoT repeatedly invokes a transformer backbone without decoding intermediate representations into tokens, while maintaining reasoning entirely within the latent space. Following this paradigm...

    2025

  53. [53]

    Memory footprint is evaluated separately by measuring peak GPU memory usage during execution, reported in GB

    Throughput is computed based on the theoretical FLOPs 𝐵𝑇𝑑 3 of solving 𝐵-many batched 𝑑-dimensional LQR problem with horizon 𝑇 via Riccati iteration and normalized by wall-clock latency. Memory footprint is evaluated separately by measuring peak GPU memory usage during execution, reported in GB. For memory benchmarking, we fix either 𝐵= 1024(when varying ...

    2020