Parallelize the Training, Keep the Learning Serial

TL;DR for operators

GradLev addresses a specific engineering problem in adaptive language models: updating temporary model weights after every token gives fine-grained adaptation, but apparently forces training to replay those updates sequentially.

The paper shows that this dependency can be moved rather than simply accepted. If the gradient signal needed for each token-level update is available, the sequence of temporary weight changes becomes a structured recurrence that can be reconstructed with parallel scans. GradLev trains an auxiliary network to predict those gradient signals, then checks them against gradients reconstructed by the primary model.¹

The formal result is strong within its assumptions: exact consistency uniquely recovers the same learner that would have performed ordinary token-by-token test-time backpropagation. A synthetic test reproduces serial outputs, internal gradients, and weight gradients to roughly $10^{-16}$ numerical precision.

The deployment story is different. GradLev does not make inference-time learning parallel across tokens. Deployment still performs token-by-token forward passes, backward passes, and weight updates. The auxiliary predictor is discarded.

For product teams, the relevant question is therefore whether extra training machinery is worth paying for to preserve fine-grained online adaptation. The current 300M experiment is encouraging but not decisive: likelihood and perplexity improve slightly, while downstream accuracy differences remain uncertain.

The bottleneck appears before deployment

Suppose a language model should learn something immediately from the current document: a newly introduced name, an unusual convention, or a local distribution shift. One option is to maintain temporary parameters that change as tokens arrive while leaving the model’s persistent learned parameters untouched.

Updating those temporary parameters after every token creates an attractive form of local adaptation. It also creates a dependency chain. Token $t+1$ uses weights changed by token $t$, whose update depended on earlier computation. Training such a system appears to require stepping through the sequence in order.

Coarser schemes can reduce that burden by updating only between chunks. But that changes the learner. The operational question is whether the per-token learner itself can be preserved without forcing its training procedure to unroll every update serially.

GradLev, introduced by Bo Liu and Qiang Liu,¹ answers this by separating the information required to define the updates from the sequential process that normally produces that information.

Once the gradient signals are known, the weight trajectory becomes parallelizable

Consider an adaptive linear map whose temporary weight at token $t$ is $W_t$, initialized from persistent weight $W_0$. GradLev uses a centered update of the form

Here $x_t$ is the layer input, $\eta_t$ controls the update size, $\alpha_t$ controls retention of previous adaptation, and $g_t$ is the gradient of the token loss with respect to that layer’s output.

That last quantity is the missing ingredient. The paper calls it a costate.

Once the inputs, gates, and costates are specified, the temporary weight correction can be written as a discounted sum of earlier rank-one gradient updates. The recurrence is affine, so partial results can be combined using an associative prefix scan rather than evaluated strictly one token at a time.

This changes the computational problem. GradLev is not trying to parallelize ordinary backpropagation through a learner whose future gradients are still unknown. It predicts the costates first, making the proposed sequence of weight updates explicit enough to reconstruct in parallel.

A causal auxiliary Transformer, $Q$, supplies those predictions across the sequence. Given them, the primary model $P$ can construct its proposed temporary-weight trajectory and its forward activations using scans.

The predictor is constrained by the learner it is supposed to reproduce

A predicted gradient signal could otherwise become merely a convenient approximation. GradLev instead makes the primary model reconstruct the costates implied by the proposed trajectory.

Training combines the normal token loss with a consistency term:

The predicted costate $\widehat{g}_t$ is pushed toward the reconstructed costate $\widetilde{g}_t$ produced by backpropagating through $P$ on the proposed trajectory.

The paper’s central theoretical result explains why this target is meaningful. Because the reconstructed costate for a token depends only on earlier proposed costates, the resulting causal mapping has a unique fixed point. Under the theorem’s deterministic, matched-update assumptions,

where $G^\star$ is the costate trajectory of the ordinary serial learner.

At exact consistency, GradLev therefore recovers not merely similar outputs but the serial learner’s temporary weights, activations, and test-time gradients in exact arithmetic.

The synthetic test checks the mechanism, not model quality

The paper first tests whether the parallel reconstruction actually reproduces serial test-time training when exact costates are supplied.

In a two-layer synthetic MLP, maximum discrepancies from serial autograd are $1.18\times10^{-16}$ for outputs, $2.22\times10^{-16}$ for hidden costates, and $6.66\times10^{-16}$ for the first matrix’s gradients. These are rounding-scale differences.

A control experiment replaces the transpose of the adapted weight with the transpose of the persistent weight during backward reconstruction. Forward outputs remain unchanged, but hidden-costate and first-matrix-gradient discrepancies jump to 0.3757 and 0.6261.

That control is an implementation-mechanism check. It shows that exact reverse reconstruction depends on reading the same adapted weight history in the backward direction; it is not evidence of broader language-model performance.

The 300M results are positive point estimates, not decisive downstream gains

The practical experiment uses a 21-layer model of roughly 300M parameters and adapts only the final MLP. Static, four-chunk test-time training, and GradLev all process 29.47 billion training tokens.

The likelihood-oriented metrics consistently favor GradLev in this experiment, although the absolute differences are small.

Downstream accuracy is less conclusive. The paired 95% interval for GradLev’s mean-accuracy difference versus Chunk-TTT is $[-0.454, 0.808]$ percentage points; versus Static it is $[-0.421, 0.862]$. Both include zero.

GradLev also pays more during training. Its deployed non-embedding parameter count is about 264.3M, close to the baselines, because $Q$ is discarded. During training, however, the model includes roughly 294.7M parameters. The parallelization removes one kind of sequential dependency; it does not make the additional predictor free.

Production value depends on scaling the gap between theory and approximation

What the paper directly establishes: the serial token-level learner can be represented through a parallel training construction when the required costates are supplied, and exact consistency recovers that learner uniquely under the stated assumptions. The synthetic diagnostic strongly verifies those mechanics.

Cognaptus inference: if learned costate prediction remains tractable as adaptation reaches more layers and larger models, this could make fine-grained online learning easier to train without retreating to chunk-level updates. Products dealing with strong document-local structure or within-context distribution shifts are plausible beneficiaries because the temporary state can change after each token while persistent model parameters remain fixed.

What remains unresolved: exact recovery is an exact-consistency result. Appendix D shows that a small consistency residual can still correspond to a large trajectory error when the causal mapping amplifies errors. Without an additional contraction or stability condition, low consistency loss alone is not a certificate that deployment dynamics are close to the intended serial trajectory.

The empirical boundary is equally concrete. Only the final MLP is adapted, the study covers one approximately 300M scale, the gains are modest, and downstream intervals do not separate GradLev from the baselines. Extending adaptation through deeper networks may also make costate prediction harder.

GradLev therefore changes the architecture of the training problem more clearly than it settles the product case. It shows that per-token test-time learning need not imply serial training across tokens. The remaining question is whether that construction stays stable and economically worthwhile when the adaptive part of the model becomes large enough to matter.

Cognaptus: Automate the Present, Incubate the Future.