Nonlinear memory stops approximating by blocks: ParaTTT turns test-time training layers into parallel refinement, and its 340M model trains faster than TTT-MLP
Synopsis
Test-time training layers replace the recurrent state with a small neural network trained online by gradient descent, which is more expressive than a linear memory but forces sequential computation because each token's gradient depends on every earlier update. ParaTTT instead iteratively refines an estimate of the hidden-layer pre-activation gradients, with each refinement built from a forward, a readout and a backward chunkwise-parallel pass, and the exact online recurrence is a fixed point of that iteration. At 4K context, two refinements reach the error of TTT-MLP's block approximation in less time, and 340M and 1.3B models match or exceed Gated DeltaNet and Mamba-2 on perplexity and recall.
Interpretation
A nonlinear fast-weight memory can now be evaluated chunkwise-parallel like a linear recurrence, instead of sharing one set of weights across a block of tokens to buy parallelism. TTT-MLP previously froze gradient evaluations per block: larger blocks cut sequential depth but approximated token-by-token learning more coarsely, while smaller blocks were more faithful but needed more sequential steps. At 4K context on one fixed random inner MLP, two ParaTTT refinements match the output error of TTT-MLP at its default block size of 16 in less time and approach the bf16 noise floor.
The method swaps the hard-to-parallel inner weight trajectory for a guess of the hidden-layer pre-activation gradients and repeatedly corrects that guess. Once this gradient guess is fixed, the nonlinear weight update becomes a linear recurrence with known coefficients that existing chunkwise-parallel kernels can evaluate. Each refinement is a forward scan, an exact gated-delta-rule readout scan and a backward scan; on the trained models, with the exact token-by-token recurrence as reference, next-token top-1 agreement is clearly low with no iteration, rises substantially with one, and moves closer to the bf16 noise floor with two.
Perspective
The demonstrated scope is 24-layer language models at 340M and 1.3B parameters trained on SlimPajama for 15.7B and 100B tokens respectively, plus single-layer kernel timings and an error-versus-time comparison on one fixed random inner MLP.
The convergence theorem covers only the two-layer variant with soft clipping in place of the unit-ball projection, a linear-attention hidden update and an uncentered readout, excluding the hard projection, the hidden delta rule and centering; the practical models train with two iterations, and decoding performs no iteration because a single new token is one exact gradient step from the cached weights.