GradLev: Token-Parallel Test-Time Training Via Costate Prediction
cs.LG, cs.AI
Submitted: 2026-09-28
Updated: 2026-09-29
Comments: Test-Time Training; Language Modeling
License: http://creativecommons.org/licenses/by/4.0/
The gist: Test-time training (TTT) allows a model to improve its predictions at inference time by updating weights after every observed token.
Abstract
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.
Sources
- ATLAS: Learning to Optimally Memorize the Context at Test Time
- It's All Connected: A Journey Through Test-Time Memorization, Attentional Bias, Retention, and Online Optimization
- Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge
- Nemotron-CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-training
- Latent Recurrent Transformer: Architecture Exploration, Training Strategies, and Scaling Behavior
- Pretraining Recurrent Networks without Recurrence
- Mamba-3: Improved Sequence Modeling using State Space Principles
- Maglev: Sliding Recurrent Memory
- End-to-End Test-Time Training for Long Context
- Next-Latent Prediction Transformers Learn Compact World Models
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks