ParaTTT: Parallelizing Test-Time Training Layers with Deep Memories

Noel Loo; Neehal Tumma; T. Konstantin Rusch; Mathias Lechner; Alexander Amini; Daniela Rus

cs.LGOct 28, 2026

ParaTTT: Parallelizing Test-Time Training Layers with Deep Memories

Authors: Noel Loo, Neehal Tumma, T. Konstantin Rusch, Mathias Lechner, Alexander Amini, Daniela Rus

Organizations: MIT CSAIL · Liquid AI · Max Planck Institute for Intelligent Systems

Abstract

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\times$ faster than TTT-MLP with the block approximation at matched error, while the full 340M model trains $4.9\times$ 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

Panel of Figure 1 — Figure 1: Parallelism and approximation in nonlinear TTT. Each cell represents a hidden pre-activation gradient $\delta\mathbf{z}_{t}(\mathbf{W})=\partial\ell_{t}/\partial\mathbf{z}_{t}|_{\mathbf{W}}$ , with the layer index omitted. Columns index tokens; downward steps mark computational dependencies, not approximation error or elapsed time. (a) Exact gradients $\delta\mathbf{z}_{t}^{\star}$ use the latest online weights $\bar{\mathbf{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 $\mathbf{W}_{s}^{\mathrm{blk}}$ for $s<t\leq 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.

Panel of Figure 2 — 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.

340M, 24 layers, 15.7B tokens
		Perplexity $\downarrow$		Zero-shot acc. $\uparrow$							Recall $\uparrow$				NIAH $\uparrow$
Model	$N$	Wiki	Lamb	PIQA	Hella	ARC-e	Wino	ARC-c	BoolQ	SciQ	SWDE	FDA	TQA	DROP	4K	8K	16K	32K
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 $\times$ 128	$128\times 128$ per head	conv 4, output gate
Mamba-2	342M	8 heads $\times$ 128, expand 1, 4 groups	$128\times 128$ per head	scalar decay
Transformer	341M	8 heads $\times$ 128, RoPE	KV cache	YaRN for RULER
TTT-MLP	349M	16 heads $\times$ 64, block 16	MLP $64\to 256\to 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

d_{\text{head}}

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

\eta_{t}=\operatorname{sigmoid}(\mathbf{w}_{\eta}^{\top}\mathbf{x}_{t}+b_{\eta})

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.

340M, 24 layers, 15.7B tokens
Model	4K	8K	16K	32K
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 $\downarrow$		Zero-shot acc. $\uparrow$
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{\times}10^{4}$	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 $\downarrow$		Zero-shot acc. $\uparrow$							Recall $\uparrow$				NIAH $\uparrow$				LongBench $\uparrow$
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.

	Layer forward+backward (ms) at $T$					Training, 8 GPUs
Variant	1K	2K	4K	8K	16K	tokens/s ( $\times 10^{3}$ )	GiB/GPU
Readout parameterization (delta-rule body, optimized kernels)
Zero-init readout	7.3	13.1	26.2	52.2	104.1	42.7	131
Learned-init readout, uncentered	8.3	14.8	29.4	58.8	116.7	38.7	136
Centered readout (PTTT)	8.5	15.1	29.8	59.8	118.8	38.0	145
Hidden-layer update (centered readout, reference implementation)

Appendix

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

Test-time training adapts model parameters or memory during deployment using each input or stream to improve performance under changing tasks and data.

See all 105 papers

DeltaTTT: Layerwise Optimization for Nonlinear Recurrent Memory
Oct 6, 2026Yining Li, Dongchen Han, Jie Fu +1Deep Learning Optimization Neural Network Optimization
Qiuzhen College, Tsinghua University · LeapLab, Tsinghua University · IQuest Research
Universal Test-Time Training
Oct 4, 2026Zefan Cai, Qinzhe Hu, Ziqiao Ma +2Memory-Augmented Language Models Recurrent Transformers
University of Wisconsin–Madison · University of Michigan–Ann Arbor · Adobe
GradLev: Token-Parallel Test-Time Training Via Costate Prediction
Sep 28, 2026Bo Liu, Qiang LiuGradient Descent Test-Time Training
The University of Texas at Austin

Oct 6, 2026J/K move · Enter open · S save

DeltaTTT: Layerwise Optimization for Nonlinear Recurrent Memory

Yining Li, Dongchen Han, Jie Fu, Gao Huang

Qiuzhen College, Tsinghua University · LeapLab, Tsinghua University · IQuest Research

Deep Learning Optimization Neural Network Optimization Test-Time Training Language Modeling

Open paper AskPDF arXiv

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.

Abstract

Test-Time Training

Explore similar work

DeltaTTT: Layerwise Optimization for Nonlinear Recurrent Memory

Universal Test-Time Training

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