Parallel method lets ai update itself token by token during use
GradLev: Token-Parallel Test-Time Training Via Costate Prediction
Machine LearningArtificial Intelligence
Summary
Sometimes AI models improve themselves while they make predictions, but this usually happens step-by-step, which is slow. The authors found a way to do these improvements for many parts of a text all at once by predicting helpful information in advance. Their method uses a helper network to guess needed values so the model can update weights in parallel instead of one by one. After training, the helper is removed, leaving the model to update naturally as it reads tokens.
What this means in practice
- •For machine learning engineers: Deploy models that can quickly adapt to new input by updating weights in parallel during token-by-token processing.
- •For natural language processing developers: Accelerate inference-time fine-tuning of language models by parallelizing gradient-based updates across input tokens.
Authors
Bo Liu, Qiang Liu
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.