Parallelizing nonlinear memory layers speeds up online training for language models
ParaTTT: Parallelizing Test-Time Training Layers with Deep Memories
Machine Learning
Summary
Some AI models learn while they are reading or processing information, which can make them smarter in real time. This learning is called test-time training and usually uses a type of memory that updates slowly and one step at a time, making it hard to do quickly. The authors came up with a new way to update these memories in parallel, or all at once, which is much faster and still accurate. Their method, ParaTTT, makes big language models train faster and perform well on language tasks, matching or beating previous approaches.
What this means in practice
- •For machine learning engineers: Speed up training of language models that adapt while processing long text by using parallel nonlinear memory updates.
- •For speech recognition developers: Improve real-time adaptation in speech models that update their memory during inference for better accuracy on diverse speakers.
Authors
Noel Loo, Neehal Tumma, T. Konstantin Rusch, Mathias Lechner, Alexander Amini, Daniela Rus
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.