Test-time training (TTT) layers replace the linear fast-weight memory of modern recurrent architectures with a small neural network trained online by gradient descent. This nonlinear memory is more expressive than the linear memories of Gated DeltaNet and Mamba-2, but is difficult to parallelize: existing implementations such as TTT-MLP rely on sequential, block-wise weight updates that run several times slower than chunkwise-parallel linear recurrences. Moreover, their block-wise execution only approximates the token-by-token online recurrence, whereas linear recurrences admit an exact chunkwise-parallel evaluation. We introduce ParaTTT (PTTT), an algorithm for parallelizing arbitrarily deep nonlinear memories by iteratively refining estimates of the inner-loop gradient trajectory. Starting from a guess of the hidden-layer gradients, each refinement consists of three chunkwise-parallel passes: a forward pass that evaluates the resulting feature trajectory, an exact gated-delta-rule pass that computes the readout trajectory and its feature gradients, and a backward pass that propagates these gradients into a new estimate. The exact online gradient descent recurrence is a fixed point of the iteration, while decoding requires no refinement iterations. We prove geometric convergence for a two-layer variant under sufficient contraction conditions. In practice, two iterations suffice. With two refinement steps at 4K context, PTTT evaluates the inner-loop trajectory 32× faster than TTT-MLP with the block approximation at matched error, while the full 340M model trains 4.9× faster. Models trained at 340M and 1.3B parameters match or exceed Gated DeltaNet and Mamba-2 and remain competitive with same-size Transformers on language-modeling perplexity and zero-shot accuracy, with better recall than the linear recurrences.
Figures & tables
Figure 1: Parallelism and approximation in nonlinear TTT. Each cell represents a hidden pre-activation gradient δzt(W)=∂ℓt/∂zt∣W , with the layer index omitted. Columns index tokens; downward steps mark computational dependencies, not approximation error or elapsed time. (a) Exact gradients δzt⋆ use the latest online weights Wˉt−1 , including decay. (b) Block-wise TTT, as used by TTT-MLP ( Sun et al., 2025 ) and Titans ( Behrouz et al., 2025b ) , evaluates gradients at shared block-entry weights Wsblk for s<t≤s+b , omitting earlier within-block updates. (c) ParaTTT refines the full gradient trajectory using three chunkwise-parallel passes per round. The block size b and refinement count N control the respective approximations; the exact trajectory is a fixed point of the ParaTTT iteration.
Figure 2: One iteration of ParaTTT; the recurrences are evaluated in chunkwise-parallel form over all T tokens. (1) A parallel forward pass evaluates the hidden layers from the current estimate of their pre-activation gradients. (2) The readout runs exactly on the resulting features and DeltaRuleGrad returns each token’s feature gradient. (3) A parallel backward pass propagates it through the hidden layers using shifted reads of the transposed recurrences (see Step 3 in Section 3.2 for details) and updates the estimate. (4) After N iterations a final forward pass and readout produce the output.
Perplexity ↓
Zero-shot acc. ↑
Recall ↑
NIAH ↑
Model
N
Wiki
Lamb
PIQA
Hella
ARC-e
Wino
ARC-c
BoolQ
SciQ
SWDE
FDA
TQA
DROP
4K
8K
16K
32K
340M, 24 layers, 15.7B tokens
GDN
–
26.57
29.66
65.2
36.5
47.1
51.5
24.1
56.6
78.4
32.5
17.8
35.7
15.3
25.9
17.0
13.7
9.3
Mamba2
–
26.99
45.87
65.7
35.1
45.5
51.9
23.4
57.8
75.2
25.8
5.7
30.6
12.6
8.1
3.0
1.1
1.2
Transformer
–
25.34
39.25
64.3
35.8
46.3
52.8
22.4
60.8
76.7
54.6
28.4
40.5
14.5
38.6
23.2
16.3
10.5
TTT-MLP
–
26.12
40.35
65.2
35.1
46.3
50.8
23.3
58.5
77.2
33.1
16.2
35.2
13.9
31.5
20.4
13.0
12.4
Table 1: Main results at 340M and 1.3B. Recall follows the generation protocol of Arora et al. (2024a) ; NIAH is the RULER needle average at the given length. Best per column and size in bold, second best underlined. A couple of notes: (1) LaCT fast-weight branch only is an ablation of the published hybrid. (2) The zero-init readout removes both centering and the learned initialization; Table 6 isolates their effects.
Figure 3: (a,b) Perplexity along 698 held-out SlimPajama documents at 340M and 1.3B. (c,d) LAMBADA perplexity and SWDE accuracy of the 340M ParaTTT models against the number of iterations at evaluation; filled marker: training N ; dotted line: Gated DeltaNet; off-chart loop-off values are printed at the edge. The uncentered model uses a zero-initialized readout.
Figure 4: Efficiency. Forward+backward time of one token-mixing layer against (a) sequence length and (b) head dimension; (c) training throughput of the full 340M models on eight MI325X GPUs. The fused TTT-MLP implementation used in our benchmarks supports head dimension 64; hollow markers show a compiled scan at larger dimensions.
Figure 5: Convergence to the sequential recurrence. (a) One random inner MLP at T=4096 : relative output error against exact online GD versus the time to compute one 340M layer’s recurrence, for ParaTTT at N=1 – 8 and TTT-MLP’s block form at b=4 – 256 (filled: fused kernel; hollow: compiled scan); the arrow marks equal error to N=2 . (b,c) Trained models on real text: top-1 agreement and total-variation distance between the prefill at N iterations and the exact recurrence.
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
Model
Params
Mixing layer
Memory / state
Notes
340M, 24 layers, d=1024
Gated DeltaNet
342M
8 heads × 128
128×128 per head
conv 4, output gate
Mamba-2
342M
8 heads × 128, expand 1, 4 groups
128×128 per head
scalar decay
Transformer
341M
8 heads × 128, RoPE
KV cache
YaRN for RULER
TTT-MLP
349M
16 heads × 64, block 16
MLP 64→256→64 per head
inner LR base 0.1
LaCT fast-weight branch only
340M
4 heads, chunk 2048
SwiGLU fast weights
Muon updates, conv 4
Appendix
Table 2: Models of Table 1 . Parameter counts include embeddings. The ParaTTT memory has a hidden layer of width dhead with SiLU, a bias, the unit-ball projection on the features, and a learned per-head gain on the fast pre-activation; its learning rates are ηt=sigmoid(wη⊤xt+bη) per head and layer, its decay is the Gated DeltaNet gate shared by both layers.
Figure 6: Synthetic nonlinear-state tasks (two layers, width 256 except width 512 for MQAR with 64 pairs; mean over three seeds at the best learning rate). (a) Cumulative parity against sequence length; the dotted line is chance. (b) k -hop pointer following at T=256 . (c) Multi-query associative recall for increasing sequence length and number of key-value pairs.
Model
Metric
N=0
1
2
3
5
7
10
340M, trained at N=2
WikiText-103
61.3
25.76
25.46
25.54
25.67
25.75
25.87
LAMBADA
332
28.43
26.93
27.37
28.21
28.40
28.60
SWDE
7.4
39.3
41.9
40.8
39.4
39.6
38.7
340M, trained at N=5
WikiText-103
61.7
26.29
25.63
25.51
25.47
25.49
25.52
LAMBADA
375
31.60
29.31
28.69
28.44
28.52
28.72
SWDE
5.0
33.8
43.1
44.0
45.0
44.2
43.8
Appendix
Table 3: Perplexity (WikiText-103, LAMBADA) and SWDE accuracy of the ParaTTT models as a function of the number of iterations at evaluation. Bold: best N for the centered models.
Model
4K
8K
16K
32K
340M, 24 layers, 15.7B tokens
GDN
11.9
12.1
12.3
12.2
Mamba2
11.8
11.9
11.9
11.9
Transformer
10.9
11.5
10.9
10.2
TTT-MLP
9.8
9.0
8.1
8.1
LaCT fast-weight branch only
4.9
4.5
4.2
4.2
Appendix
Table 4: LongBench: average score over the 16 English tasks (100 samples per task) at context lengths 4K–32K, for the models of Table 1 . Best per column and size in bold, second best underlined. All models trained at 4K; ParaTTT is within a point of Gated DeltaNet at every length and scale, and attention leads at 1.3B.
Perplexity ↓
Zero-shot acc. ↑
Model
Wiki
Lamb
PIQA
Hella
ARC-e
Wino
ARC-c
BoolQ
SciQ
LaCT with SWA-2048 (as published)
25.51
33.3
66.2
35.8
45.1
51.7
23.6
61.1
79.3
LaCT fast-weight branch only ( Table 1 )
463.2
9.4×104
53.7
25.0
28.8
50.1
24.1
60.1
21.2
PTTT (ours) , no attention
25.46
26.9
65.3
37.7
48.6
51.2
24.1
59.8
79.0
Appendix
Table 5: LaCT at 340M with and without its sliding-window attention branch, same recipe and data as Table 1 ; recall for the fast-weight-branch-only ablation is in Table 1 .
Perplexity ↓
Zero-shot acc. ↑
Recall ↑
NIAH ↑
LongBench ↑
Readout
Wiki
Lamb
PIQA
Hella
ARC-e
Wino
ARC-c
BoolQ
SciQ
SWDE
FDA
TQA
DROP
4K
8K
16K
32K
4K
8K
16K
32K
Gated DeltaNet
26.57
29.66
65.2
36.5
47.1
51.5
24.1
56.6
78.4
32.5
17.8
35.7
15.3
25.9
17.0
13.7
9.3
11.9
12.1
12.3
12.2
Zero-init
26.06
32.62
64.5
36.8
48.4
50.6
25.3
58.2
78.9
36.6
19.6
37.0
14.3
29.0
21.1
15.0
12.2
12.6
12.6
12.5
12.5
Learned init, uncentered
26.13
29.06
65.8
36.7
46.1
52.6
24.4
59.7
80.4
37.7
19.2
37.1
13.5
32.5
19.2
14.4
14.4
11.6
11.2
11.2
11.4
Centered (PTTT)
25.46
26.93
65.3
37.7
48.6
51.2
24.1
59.8
79.0
41.9
20.3
39.6
14.4
32.6
20.7
15.6
16.0
12.4
12.5
12.4
12.1
Appendix
Table 6: Centering ablation at 340M. Three parameterizations of the readout layer, trained with the same recipe, initialization and data order; the gated DeltaNet baseline is shown for reference. Columns as in Table 1 ; best of the three PTTT variants per column in bold.
Table 7: Cost of the hidden-layer update and of the readout parameterization at 340M, N=2 . Left: forward+backward time of one token-mixing layer (bf16, batch 8, d=1024 , 8 heads of dimension 128, one MI325X), the setting of Fig. 4 (a). Right: training throughput and peak memory of the full model on one node of eight MI325X GPUs with the recipe of Appendix D . The readout rows use the optimized kernels of Appendix G ; the two bodies are timed in the reference implementation, which supports both updates, so the body rows are comparable with each other but slower than the readout rows.
Figure 7: Relative error against exact online GD versus the time to compute the forward recurrence of one 340M layer (batch 1), at five sequence lengths; same setting and legend as Fig. 5 (a).
Test-time training adapts model parameters or memory during deployment using each input or stream to improve performance under changing tasks and data.
Sequential test-time training adapts a memory network through successive updates, each computing an inner-loop gradient based on the network's previous state. Intuitively, this state dependence should allow each update to account for what the memory has already learned and better incorporate new information. However, we find that this expected advantage does not consistently materialize in nonlinear memories: a fixed-base parallel TTT baseline outperforms its serial counterpart. Our exploratory experiments point to a key underlying difficulty: nonlinear memories can be harder to optimize than linear ones within a single pass over the sequence. To alleviate this optimization difficulty, we introduce DeltaTTT, which replaces joint inner-loop optimization of a two-layer memory network with layerwise learning. Each layer is assigned a local prediction target and updated through a state-dependent delta rule. This formulation retains a nonlinear readout while enabling chunkwise parallel computation. Experiments on DeltaNet and LaCT backbones show improvements in language modeling and retrieval over their recurrent baselines.
Yining Li, Dongchen Han, Jie Fu +1
Qiuzhen College, Tsinghua University · LeapLab, Tsinghua University · IQuest Research
Recent Test-Time Training (TTT) architectures compress context into fast weights that are updated online and queried as memory. Existing TTT designs keep this memory private to each layer: it recurs only over time, and depth merely indexes L separate memories. We argue that memory ownership need not be tied to depth, and introduce Universal Test-Time Training (uTTT), in which all layers read and write one shared memory while retaining layer-specific backbone parameters. The shared memory thus recurs over two dimensions, time and depth, with chunks and layers as their units: a write by a deep layer in one chunk can be read by a shallow layer in the next. We instantiate this idea as uTTT-MoE and uTTT-Dense. uTTT-MoE routes each token head to a few experts in a pool shared by all layers; uTTT-Dense applies the whole shared memory at every layer without routing. In language modeling, uTTT-MoE reaches 15.5 and 27.9 RULER accuracy at 124M and 760M, 2.6 and 2.1 points above its layer-private counterpart at equal state and active compute, the highest among tested bounded-state models, with per-token loss matching or beating full attention. In novel view synthesis, sharing at fixed per-layer compute gains 0.92 dB in view-23 object PSNR in routed models and 0.76 dB in dense models.
Zefan Cai, Qinzhe Hu, Ziqiao Ma +2
University of Wisconsin–Madison · University of Michigan–Ann Arbor · Adobe
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.