PaperlayerView exact revision on arXiv

Source-linked AI summary

GradLev: Token-Parallel Test-Time Training Via Costate Prediction

Bo Liu, Qiang Liu

arXiv:2609.34174v2cs.LGcs.AI

TL;DR

Token-level test-time training uses each observed token to update weights, but sequential gradient writes make parallel training difficult. GradLev predicts costates and uses scans with a consistency objective to train the learner; exact consistency recovers the serial online trajectory, while reported empirical gains are modest and larger-scale effectiveness remains unestablished.

  • Problem

    Token-level test-time training adapts weights after observed tokens, but sequential gradient writes make sequence-wide parallel training difficult.

    Source: §1 · pp. 1–2

  • Method

    GradLev predicts costates across tokens and uses parallel scans for forward evaluation and backward reconstruction, with a consistency loss aligning proposals to reconstructed targets.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    At exact consistency, parallel matrices, activations, and test-time gradients coincide with the serial learner in exact arithmetic.

    Source: §3 · p. 3

  • Takeaways & Limitations

    At deployment, GradLev discards the predictor and updates the primary learner online with exact per-token gradients.

    Source: §1 · p. 2

    The preliminary experiments show modest empirical gains and do not establish effectiveness at larger scales.

    Source: §6 · p. 6

Abstract

from arXiv · show

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.

1 INTRODUCTION

Test-time training adapts temporary weights from each observed token, but token-by-token backpropagation creates a serial dependency that remains an open challenge for sequence-wide parallel training. GradLev predicts costates across tokens and uses parallel scans to train the learner, with a consistency loss aligning predictions to reconstructed gradients.

1 INTRODUCTION

  • Each incoming token is both a prediction target and an immediate supervisory signal, allowing temporary weights to adapt before the next prediction.The paper connects this online adaptation principle to tracking long-range document context, bursty entity bindings, and local distribution shifts.Source: §1 · p. 1
  • Token-level test-time backpropagation serializes prediction, gradient computation, and weight updates, while chunked updates sacrifice immediate reactivity and simpler recurrences constrain model expressiveness.Training an expressive model with a full forward pass, backward pass, and weight update at every token while retaining sequence-wide parallel training remains an open challenge.Source: §1 · pp. 1–2
  • GradLev uses a causal auxiliary network to predict costates across tokens, enabling parallel scans for the learner’s forward states, adapted predictions, and transposed backward passes.Given predicted costates and layer inputs, online gradient descent becomes an affine matrix recurrence.Source: §1 · p. 2
  • A consistency loss aligns predicted costates with targets that the primary learner reconstructs by retracing backpropagation through scanned weights.At zero consistency error, the parallel training trajectory matches the causal online trajectory; deployment discards the predictor and updates the learner online.Source: §1 · p. 2

2 TEST-TIME BACKPROPAGATION

Test-time backpropagation scores each next token using parameters adapted from the observed prefix. Its update rule operates on fast weights, while meta-learning trains the initialization and update rule that remain part of the model’s persistent parameters.

2 TEST-TIME BACKPROPAGATION

  • At each step, θ_t predicts the next token from the observed prefix, and the token’s negative log-probability defines its loss.The prefix and historical cache are implicit in the notation.Source: §2 · p. 2
  • Each target is scored before the update, and its test-time gradient holds earlier updates and cached activations fixed.The outer objective takes an expectation over sequences.Source: §2 · p. 2
  • Meta-learning trains the initialization and update rule, while selected fast weights adapt online and slow weights remain fixed at deployment.Slow weights, including learned update-rule coefficients, are trained by the outer objective.Source: §2 · p. 2

3 GRADLEV

GradLev predicts costates in parallel, then uses scans to construct adapted forward states and reconstruct gradients backward. Its causal consistency theorem establishes exact recovery of the serial costate trajectory, subject to stated assumptions and a distinction between test-time gradients and full meta-gradients.

3 GRADLEV

  • Given inputs, costates, and gates, gradient descent becomes an exact parallel scan whose writes support both forward evaluation and backward reconstruction.The predictor supplies costate vectors; the learner reuses the same historical writes for forward and transpose scans.Source: §3 · p. 2
  • The causal predictor proposes costates across tokens, and an exclusive scan uses proposals from earlier positions to compute adapted weights for each token.Proposals may use the current target because their writes affect only weights at later positions; the update is initialized at zero and read before each write.Source: §3 · p. 3
  • Backward reconstruction passes output costates through the scanned weights, using a transpose scan over the same historical writes as the forward pass.The reverse pass holds incoming weights and historical caches fixed; repeating it across layers yields each layer’s costate on the proposed trajectory.Source: §3 · p. 3
  • Theorem 1 gives a unique fixed point because each reconstructed target depends only on earlier proposed costates.At that solution, parallel and serial matrices, activations, and test-time gradients coincide in exact arithmetic.Source: §3 · p. 3
  • Exact consistency guarantees exact test-time gradients, not the full meta-gradient; first-order MAML slow-weight recovery additionally requires fixed update coefficients and detached gradient writes.The implementation also learns gates and retains their derivatives.Source: §3 · p. 3
  • GradLev’s attention reverse pass treats past key-value cache entries as fixed constants while reading them unchanged.This constraint applies when adapting GradLev to architectures with attention.Source: §3 · p. 4

4 RELATED WORK

Related work places GradLev among online adaptation, fast-weight learning, associative recurrences, and methods for parallelizing nonlinear recurrence. Its distinguishing approach is to predict costates for scan-based reconstruction while retaining token-level updates at deployment.

Fast weights and learning to learn

  • GradLev learns an initialization and explicit gradient-descent gates, then differentiates next-token loss at deployment to update selected weights.This differs from approaches that interpret in-context learning as optimization within linear attention.Source: §4 · p. 4

Test-time training and adaptive language models

  • Unlike chunk-based TTT approaches, GradLev parallelizes training through costate prediction while preserving token-level updates at deployment.LaCT uses large update chunks for hardware efficiency and nonlinear fast-weight capacity.Source: §4 · p. 4

SSMs and memory as online learners

  • GradLev adapts selected maps using predictive loss and reconstructs their forward and reverse passes with scans, unlike memory methods with distinct regression or associative-memory objectives.Examples in the passage include Longhorn, Titans, Miras, Atlas, and MesaNet.Source: §4 · p. 4

Efficient matrix recurrences

  • GradLev reuses gated additive scans with negative costate writes rather than Gated DeltaNet’s associative-regression residual.The related-work discussion also situates these scans among efficient matrix recurrences and state-space models.Source: §4 · pp. 4–5

Parallel training of nonlinear recurrence

  • GradLev predicts costates supervised by a full reverse pass, removes the predictor at deployment, and constructs the training trajectory in three stages without iterative refinement.This contrasts with related methods that iteratively refine trajectories or train auxiliary recurrent-state predictions.Source: §4 · p. 5

5 EXPERIMENTS

The experiments test exact scan recovery for token-level MLP learning and evaluate learned costate prediction in a 300M language model adapting only its final MLP. GradLev shows modest deployment improvements over chunk-TTT and a static Transformer in that language-model setting.

5.1 RECOVERING TOKEN-LEVEL TTT IN AN MLP

  • With supplied exact costates, maximum output, hidden-costate, and first-matrix-gradient discrepancies from serial TTT are 1.18 × 10−16, 2.22 × 10−16, and 6.66 × 10−16, respectively.Errors are maxima over coordinates, tokens, and three seeds.Source: Figure 2 · p. 5, §5.1 · p. 5
  • Replacing only the adapted transpose preserves outputs but changes reverse gradients; thus, matching forward outputs alone does not establish matching reverse gradients.Errors are maxima over coordinates, tokens, and three seeds.Source: Figure 2 · p. 5
  • The exact-recovery experiment supplies costates to one parallel construction, so it verifies numerical equivalence rather than learned costate-prediction accuracy or regression generalization.The experiment compares parallel forward and backward scans with serial TTT.Source: §5.1 · pp. 5–6

5.2 300M LANGUAGE MODELING WITH FINAL-MLP ADAPTATION

  • GradLev uses a two-block causal Transformer Q to predict costates from MLP inputs and observed target embeddings; deployment removes Q.Both adaptive models update only the final MLP, and evaluation scores before updating and resets at document or context boundaries.Source: §5.2 · p. 6
  • GradLev’s final NLL is lower by 0.001856 than chunk-TTT and by 0.002617 than static; mean accuracy is 50.99%, versus 50.81% and 50.77%, respectively.GradLev leads both controls at every shared checkpoint from step 65,536 onward.Source: Table 1 · p. 6, Figure 3 · p. 6, §5.2 · p. 6

6 LIMITATIONS AND FUTURE WORK

This preliminary study reports modest empirical gains and does not establish effectiveness at larger scales. Broader benefits across model sizes, adaptation depths, and tasks, as well as several architectural choices, remain open.

6 LIMITATIONS AND FUTURE WORK

  • The empirical gains are modest, and the experiments do not establish effectiveness at larger scales or consistent benefits across model sizes, adaptation depths, and tasks.The work characterizes itself as a preliminary study of costate prediction for parallelizing token-level test-time training.Source: §6 · p. 6
  • Deeper adaptation requires a scalable predictor that maintains accurate costate predictions across layers while preserving parallel-training benefits.Parameter sharing between P and Q, and which parameters to adapt, remain open design questions with accuracy, cost, and stability trade-offs.Source: §6 · p. 7

A WORKED EXAMPLE: AN ONLINE MLP

The worked example derives online MLP updates and their parallel forward and reverse computations. The forward uses costate proposals to scan adapted weights, while the reverse reconstructs targets from the scanned computation.

A WORKED EXAMPLE: AN ONLINE MLP

  • The MLP defines preactivation, hidden activation, and output from each input using two adapted weight matrices.The matrices start from W1,0 and W2,0, and the activation is elementwise.Source: A WORKED EXAMPLE: AN ONLINE MLP · p. 10
  • Retention α_t and shared write coefficient μ_t govern centered updates, with both gradients computed using pre-update weights.The write coefficient is divided by each map’s input width.Source: A WORKED EXAMPLE: AN ONLINE MLP · p. 10
  • Given proposed costates, the forward processes layers in order and computes each recurrence with an exclusive token scan.The reverse uses historical costates to transpose-scan the same matrices and supplies targets for the consistency loss.Source: A WORKED EXAMPLE: AN ONLINE MLP · p. 10
  • At deployment, true costates are written after each scored target, with λ = 1.Source: A WORKED EXAMPLE: AN ONLINE MLP · p. 10
  • The final MLP’s token loss is formed from the residual input and tied embedding/output matrix, with RMSNorm.Its output gradient is differentiated through the head, normalization, and residual before the GELU derivative and adapted transpose yield the preactivation costate.Source: A WORKED EXAMPLE: AN ONLINE MLP · p. 10

B SYNTHETIC MLP PROTOCOL

The synthetic protocol checks whether a parallel reconstruction matches serial autograd, then specifies how forward and transpose scans reuse historical writes. The numerical check uses fixed exact costates rather than a learned predictor.

B SYNTHETIC MLP PROTOCOL

  • The protocol uses seeds 17, 42, and 123, Gaussian slow-matrix initializations, and input-conditioned retention and step sizes.The vectors controlling retention and write strength have N(0, 0.25^2) entries; matrix step sizes are μ_t/4 and μ_t/7, with both matrices updated simultaneously using centered decay.Source: B SYNTHETIC MLP PROTOCOL · p. 10
  • The reconstruction uses exact serial-autograd costates and explicit causal sums in FP64, without reusing serial hidden states.Errors are maxima over coordinates, 32 token positions, and three seeds; this is a numerical consistency check without learned Q or refinement.Source: B SYNTHETIC MLP PROTOCOL · p. 11

C FORWARD AND TRANSPOSE SCAN READS

  • Forward queries use a_t = x_t and backward queries use b_t = e^g_t, while both read historical proposals.Each read excludes the current write and all writes before the most recent document or context reset.Source: C FORWARD AND TRANSPOSE SCAN READS · p. 11

D WHAT APPROXIMATE CONSISTENCY GUARANTEES

Exact consistency and stability to prediction errors are distinct guarantees. A small consistency residual does not necessarily imply a small trajectory error when amplification is large.

D WHAT APPROXIMATE CONSISTENCY GUARANTEES

  • Causality alone cannot bound trajectory error by a small consistency residual, even when the causal target map has a unique fixed point.For the example F(G0, G1) = (0, κG0), proposals (ε, κε) have residual |ε| while the error can grow with amplification.Source: D WHAT APPROXIMATE CONSISTENCY… · p. 11

E LANGUAGE-MODEL TRAINING DETAILS

The language-model experiments specify the final-MLP architecture, optimization and data handling, and a chunked-update comparator. They report separate curve and final-score evaluations, whose endpoint NLLs differ only slightly.

E LANGUAGE-MODEL TRAINING DETAILS

  • The 300M backbone uses context length 1024, hidden and MLP widths of 1024 and 4096, and a tied 32768-by-d embedding/output matrix.Q receives normalized final-MLP inputs and the next-token embedding, then produces costate proposals from two causal Transformer blocks.Source: E LANGUAGE-MODEL TRAINING DETAILS · p. 11
  • GradLev uses half-rate Q and consistency weight λ = 1, with 262,144 scored tokens per update and a 112,420-step schedule.AdamW uses peak rate 3 × 10^-4, followed by cosine decay to 10% of peak after 128 warmup steps.Source: E LANGUAGE-MODEL TRAINING DETAILS · p. 12
  • Training resets fast state at document and packed-row boundaries and holds out the final sorted shard.Source: E LANGUAGE-MODEL TRAINING DETAILS · p. 12
  • The chunk comparator uses four fixed 256-token chunks, unit retention, and a write gate initialized to 0.9.At chunk boundaries, both matrices receive fragment-summed gradients normalized by input width; only GradLev uses a prefiller and consistency objective.Source: E LANGUAGE-MODEL TRAINING DETAILS · p. 12
  • Curve evaluations use GB300 and final scores use B300 replays, with endpoint NLLs differing by at most 1.02 × 10^-4.The two evaluation sources are reported separately.Source: E LANGUAGE-MODEL TRAINING DETAILS · p. 12

F DOWNSTREAM EVALUATION

The downstream evaluation uses an eight-task accuracy suite and separate likelihood metrics, but reports mixed task outcomes and mean-accuracy intervals that include zero. Its protocol also limits evidence about repeated chunk adaptation.

F DOWNSTREAM EVALUATION

  • The benchmark equally weights accuracy across eight tasks; examples start with zero fast state, and each token is scored before updates.Candidate prompts use independent state copies, and gold labels are consulted only after scoring.Source: F DOWNSTREAM EVALUATION · p. 12
  • Paired 95% intervals for GradLev’s mean-accuracy difference are [−0.454, 0.808] percentage points vs chunk-TTT and [−0.421, 0.862] vs static; both include zero.The intervals omit training-seed variability, and validation and FineWeb-Edu scores were used during development, so neither provides fresh confirmation.Source: Table 3 · p. 13, F DOWNSTREAM EVALUATION · p. 13
  • Chunks remain 256 tokens; many benchmark requests finish before an update affects scoring, while 1024-token NLL evaluations cross boundaries.This limits evidence about repeated chunk adaptation.Source: F DOWNSTREAM EVALUATION · p. 13
  • Likelihood evaluation separates primary validation NLL, FineWeb-Edu BPB, and LAMBADA word perplexity; 1024-target blocks score every FineWeb-Edu text token once.FineWeb-Edu evaluation uses 512 documents, while LAMBADA word perplexity averages final-word losses over 5153 examples.Source: F DOWNSTREAM EVALUATION · p. 13
  • Accuracy intervals use paired bootstrap draws within each task and equally average task differences, excluding training-seed and recipe-selection uncertainty.The procedure uses 10,000 draws and percentile limits.Source: F DOWNSTREAM EVALUATION · p. 13

Related papers

Loading 2609.34174v2…