GradLev: Token-Parallel Test-Time Training Via Costate Prediction

Bo Liu; Qiang Liu

cs.LGSep 28, 2026

GradLev: Token-Parallel Test-Time Training Via Costate Prediction

Authors: Bo Liu, Qiang Liu

Organizations: The University of Texas at Austin

Abstract

Test-time training (TTT) allows a model to improve its predictions at inference time by updating weights after every observed token. However, sequential gra- dient writes make parallel training difficult. We observe that, given layer inputs and activation gradients (costates), online gradient descent admits exact parallel scans for both forward evaluation and reverse backpropagation. GradLev lever- ages this duality: a causal auxiliary network predicts costates across all tokens in parallel; associative scans compute the adapted weights and forward activations and propagate gradients backward; and the resulting gradient targets supervise the predictor via a consistency loss. Exact consistency guarantees exact recovery of the sequential online learner. At deployment, the auxiliary predictor is discarded, and the model updates natively via token-by-token forward and backward passes.

Figures & tables

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Deployed (M)	Training (M)	Tokens (B)	Tokens / param.
Static	264.2852	264.2852	29.4702	$111.5092\times$
Chunk-TTT (4)	264.2862	264.2862	29.4702	$111.5088\times$
GradLev	264.2862	294.7000	29.4702	$100.0008\times$

Appendix

Table 2: 300M final-MLP parameter and token budgets. Counts exclude the tied embedding/output matrix. Training counts include Q and determine the reported token-to-parameter ratios.

	Likelihood ( $\downarrow$ )			Zero-shot accuracy (%) ( $\uparrow$ )
Model	Valid. NLL	FineWeb BPB	LAMB. PPL	LAMB.	PIQA	Hella.	Wino.	ARC-E	ARC-C	SIQA	BoolQ	Mean
Static	2.573312	0.752708	24.661	37.69	72.52	47.84	53.99	65.49	33.45	39.87	55.35	50.77
Chunk-TTT (4)	2.572551	0.752980	24.538	37.40	72.09	47.73	53.91	65.19	31.83	38.18	60.18	50.81
GradLev	2.570695	0.751658	23.564	38.04	72.14	47.75	53.28	65.53	32.00	39.41	59.82	50.99

Appendix

Table 3: 300M final-MLP results after 29.47B training tokens. All models use the same B300 evaluation runtime. Metrics retain their original precision; bold marks the best value in each column.

Explore similar work

Modular TTT: Rethinking Test-Time Training as Composable Modules
Aug 7, 2026Bohao Tang, Zhen Qin, Yuqi Pan +3Test-Time Training Sequence Modeling
Shanghai Jiao Tong University · Shanghai Innovation Institute · ByteDance Seed
Test-Time Training with Next-Token Prediction
Jun 19, 2026Xuan Ouyang, Zefan Cai, Junjie HuTest-Time Training Next-Token Prediction
University of Wisconsin–Madison
RW-TTT: Batched Serving for Request-Owned Test-Time Training State
May 27, 2026Jian Yang, Zhizhuo Kou, Yao Tian +4Large Language Model Serving Time-To-First-Token
1The Hong Kong University of Science and Technology · 2The Chinese University of Hong Kong · 3National University of Singapore

Aug 7, 2026, cs.LGJ/K move · Enter open · S save

Modular TTT: Rethinking Test-Time Training as Composable Modules

Bohao Tang, Zhen Qin, Yuqi Pan, Zheng Li+2

Shanghai Jiao Tong University · Shanghai Innovation Institute · ByteDance Seed

Test-Time Training Sequence Modeling Batch

Open paper AskPDF arXiv

Test-time training (TTT) views sequence modeling as an online learning problem in which fast weights are updated by an internal learning rule. Despite the growing number of TTT variants, existing approaches typically hard-code each variant separately, which makes it difficult to design new TTT methods and to isolate the role of each component. To address this, we propose Modular TTT, a framework that represents the inner learner as a directed acyclic graph and exposes the fast-weight network, loss function, learning rate, weight decay, and normalization as explicit design dimensions. Modular TTT automatically composes primitive-level train-view forward, train-view backward, and causal query-view rules into the full graph-level TTT computation, including the fast-weight state transition. Using Modular TTT, we systematically ablate the components of TTT and find that small learning-rate initialization, weight decay, and a single-layer nonlinearity improve performance, while MSE and inner-product losses perform similarly. Deeper fast-weight networks and normalization tend to hurt performance because they induce excessively large activations, while residual connections and gating provide little measurable benefit. Guided by these findings, we train the best resulting variant as 410M- and 1.45B-parameter models on 100B tokens, and observe training loss and benchmark performance comparable to Gated DeltaNet.

Abstract

Explore similar work

Modular TTT: Rethinking Test-Time Training as Composable Modules

Test-Time Training with Next-Token Prediction

RW-TTT: Batched Serving for Request-Owned Test-Time Training State