ParaTTT: Parallelizing Test-Time Training Layers for Deep Memory

The Parallelization Dilemma of Test-Time Training Layers

Test-time training (TTT) layers replace the linear fast-weight memory in modern recurrent architectures with a small neural network that is trained online via gradient descent. This nonlinear memory is more expressive than the linear memory of Gated DeltaNet and Mamba-2, but it is difficult to parallelize: existing implementations such as TTT-MLP rely on sequential, chunked weight updates and run several times slower than chunk-parallel linear recurrences. Moreover, their chunked execution only approximates the token-by-token online recurrence, whereas linear recurrences can be evaluated with exact chunk parallelism.

Test-time training (TTT) is a mechanism that lets a model continue adjusting its internal parameters via gradient descent even at inference time. The core motivation is that standard forward-pass weights are frozen at inference, so the model cannot adapt to the current input stream; TTT instead places part of its "memory" in a small network that can be updated online, allowing it to evolve as the sequence progresses. This differs from "test-time augmentation" or "test-time adaptation," which typically only adjust the input or a few normalization parameters, whereas TTT directly takes gradient steps on the prediction network inside the recurrence.

Fast weights refer to weights that are not learned long-term via backpropagation but are computed on the fly from the current context during the forward pass. Linear fast-weight memory (such as the state matrices in Gated DeltaNet and Mamba-2) can be written as matrix recurrences and thus parallelized in chunks; nonlinear fast weights encapsulate the recurrence inside an MLP, which is more expressive but breaks the parallelizable linear structure.

Chunk parallelism is a key technique for efficient training of linear recurrences: a long sequence is split into blocks, each block is computed at once via parallel scan or matrix multiplication, and only a low-dimensional state is passed between blocks. The speed of Gated DeltaNet and Mamba-2 comes precisely from the fact that their linear recurrences admit exact chunk-parallel evaluation. The nonlinear updates of TTT-MLP cannot be written in this form, so gradient steps must be taken sequentially block by block, and the within-block approximation deviates from the exact token-by-token online recurrence.

Gated DeltaNet is an improved version of DeltaNet that adds gating on top of the delta rule (correcting old memory with new key-value pairs) to control memory writes and forgetting. Mamba-2 is a structured state space model whose core layer can be viewed as a gated linear attention/recurrence that supports efficient chunk parallelism. Both are representative of current linear recurrent architectures and serve as the main baselines for ParaTTT.

The Core Idea of ParaTTT

ParaTTT (PTTT) parallelizes nonlinear memory of arbitrary depth by iteratively refining an estimate of the inner-loop gradient trajectory. Starting from a guess for the hidden-layer gradients, each refinement consists of three chunk-parallel processes: a forward pass that evaluates the resulting feature trajectory; an exact gated-delta-rule process 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 loop is the fixed point of this iteration, and decoding requires no refinement iterations.

Fixed-point iteration means constructing a mapping such that the target solution is exactly a fixed point of that mapping, then repeatedly applying the mapping from an arbitrary initial value to approach the solution. ParaTTT treats the exact online gradient descent trajectory as the fixed point of the composite mapping "forward pass—readout—backward pass": if the guessed hidden-layer gradients exactly equal those produced by the exact recurrence, then the trajectory is unchanged after one refinement. Thus convergence of the iteration means that the parallel computation reproduces the result of sequential online training.

The contraction condition is a sufficient condition guaranteeing geometric convergence of the fixed-point iteration; intuitively, it requires the mapping to be contractive with respect to gradient perturbations (Lipschitz constant less than 1). The authors prove that the two-layer variant has errors decaying geometrically under sufficient contraction, which explains why two refinements suffice in practice: the initial value (such as zero gradients or the previous block's gradients) is already close to the fixed point, and contraction rapidly shrinks the residual.

Convergence and Performance

The authors prove that the two-layer variant converges geometrically under sufficient contraction. In practice, two iterations suffice. At a 4K context, PTTT with two refinement steps evaluates the inner-loop trajectory 32 times faster than TTT-MLP with chunked approximation at matched error, while full 340M model training is 4.9 times faster. Models trained at 340M and 1.3B parameters match or exceed Gated DeltaNet and Mamba-2 on language modeling perplexity and zero-shot accuracy, remain competitive with same-size Transformers, and have better recall capability than linear recurrences.

Perplexity is the standard metric for language modeling, equal to the exponential of the average negative log-likelihood—lower is better; zero-shot accuracy refers to the accuracy of a model answering directly in the task format without downstream fine-tuning, used to measure general capability. Recall capability here refers to the model's ability to retrieve and reproduce earlier tokens or facts from context, a key dimension distinguishing linear recurrences from Transformers: linear memory has limited capacity, and long-range exact recall is typically weaker than attention, and ParaTTT claims to outperform linear recurrences on this item.

340M and 1.3B are parameter scales, belonging to small-to-medium-scale controlled experiments; 4K context refers to a training/evaluation sequence length of 4096 tokens. The 32x and 4.9x figures are the speedups for inner-loop trajectory evaluation and full model training, respectively; the former measures the direct benefit of parallelization to the memory layer, while the latter is diluted by attention and other layers, hence the smaller speedup.

One-line Summary

ParaTTT uses fixed-point iteration to parallelize nonlinear test-time training layers, training several times faster without sacrificing performance.

Paragraph Summary

The article introduces ParaTTT (PTTT), a method for parallelizing nonlinear memory in test-time training (TTT) layers. TTT layers replace linear fast weights with a small network trained via online gradient descent, offering greater expressiveness, but implementations such as TTT-MLP rely on sequential chunked updates that are slow and only approximate the token-by-token recurrence. ParaTTT treats the exact online gradient descent loop as a fixed point and iteratively refines it via three chunk-parallel steps: forward evaluation of the feature trajectory, exact gated-delta-rule computation of the readout and feature gradients, and backward propagation to update the gradient estimate; decoding requires no iterations. The authors prove geometric convergence of the two-layer variant under contraction, and in practice two iterations suffice; at 4K context, inner-loop trajectory evaluation is 32 times faster than TTT-MLP, and 340M model training is 4.9 times faster, matching or exceeding Gated DeltaNet, Mamba-2, and same-size Transformers on perplexity, zero-shot accuracy, and recall.