PaperlayerView exact revision on arXiv

Source-linked AI summary

Rethinking Expressivity and Efficiency in Test-Time Training

Zeyun Zhong, Joya Chen, Manuel Martin, Frederik Diederichs, Juergen Gall, Juergen Beyerer

arXiv:2608.21308v1cs.LG

TL;DR

TTT methods face a trade-off between expressive token-wise updates and efficient chunk-wise processing for long-context modeling. E2-TTT derives a closed-form transition that exactly reproduces chunk-end states under frozen chunk-start gradients while enabling parallel chunk updates. It is competitive in language modeling, stronger in retrieval and length extrapolation, and sustains > 90% passkey accuracy at 8× the training context length.

  • Problem

    Existing TTT methods struggle to balance expressive per-token update dynamics with the hardware efficiency of chunk-wise approximations for long-context processing.

  • Method

    E2-TTT derives a closed-form chunk-parallel state transition that reproduces the chunk-end fast-weight and momentum states of a per-token recurrence under frozen chunk-start gradients.

  • Results

    E2-TTT is on par with strong sub-quadratic baselines in language modeling, outperforms them on retrieval, and sustains > 90% passkey accuracy at 8× the training context length.

  • Takeaways & Limitations

    Preserving precise temporal update dynamics supports long-context generalization while retaining the hardware efficiency of chunk-wise processing.

  • Takeaways & Limitations

    The method assumes gradients are evaluated at fixed chunk-start weights, which decouples token gradients across each chunk.

Abstract

from arXiv · show

Test-Time Training (TTT) enables long-context processing via continuous weight updates during inference, but current methods struggle to balance the expressivity of per-token update dynamics with the hardware efficiency of chunk-wise approximations. We propose E$^2$-TTT (Expressive and Efficient TTT) to bridge this gap. Under the standard approximation of taking gradients at the chunk-start weights, we derive a closed-form state transition that exactly reproduces the chunk-end fast-weight and momentum states of the per-token recurrence. This enables fully parallelized chunk-level training while preserving the temporal structure of the update rule that prior chunk-wise methods discard. We validate E$^2$-TTT by training models up to 1.3B parameters from scratch. It performs on par with previous TTT and hybrid attention baselines in language modeling while outperforming them on in-context retrieval. Its advantage is most pronounced in length extrapolation: on the standard ``Needle in a Haystack'' passkey test, it retains over 90% accuracy at $8\times$ the training context length. Meanwhile, E$^2$-TTT can match the training throughput of efficient chunk-wise methods, demonstrating that it effectively reconciles expressivity with efficiency. The code is available at https://github.com/zeyun-zhong/E2-TTT.

1 Introduction

TTT addresses long-context limitations by adapting fast weights during inference, but existing methods trade token-level expressivity for efficient chunk-wise computation. E2-TTT preserves per-token update dynamics while enabling chunk-level parallelism and shows strongest gains in retrieval and length extrapolation.

  • Frozen LLM weights leave models dependent on accumulated memory for long-context tasks and limit continuous adaptation across long-horizon or cross-task settings.
  • Linear or hybrid architectures reduce attention complexity, while TTT adapts fast-weight subnetworks online to support longer-sequence processing.
  • Token-wise TTT offers expressiveness but incurs sequential training, whereas chunk-wise TTT improves efficiency by simplifying temporal update dynamics.
  • E2-TTT uses a closed-form parallel scalar kernel with per-token learning rate, momentum, and decay while preserving each token’s effect on the chunk-end state.
  • Models up to 1.3B parameters were trained from scratch, with evaluation covering language modeling, in-context retrieval, and length extrapolation.
  • E2-TTT maintains > 90% passkey-retrieval accuracy at 8× the training context length in the Needle in a Haystack test.

2 Background: Test-Time Training

TTT treats hidden states as learnable fast weights that are updated from in-context key-value associations and then applied to queries. Mini-batch and chunk-wise variants improve parallelism, but they introduce approximations and use coarser state updates.

  • TTT represents recurrent state as learnable weights of a sequence-dependent nonlinear function, while slow model parameters remain fixed during inference.
  • Each TTT step projects inputs into queries, keys, and values; keys and values update fast weights, which then generate the output for the query.
  • Token-wise updates reconstruct values from transformed keys, then immediately apply the updated fast weights to the current query.
  • Mini-batch TTT evaluates gradients at fixed previous-batch weights, decoupling forward dependencies so gradients can be computed in parallel.
  • Parallel associative scans can compute token-level fast-weight states after gradients are calculated, but token-wise dependencies still limit hardware utilization.
  • Chunk-wise TTT stores one set of gradient, momentum, and weight states per chunk and computes the chunk’s outputs using the previous chunk’s weights.

3 Expressive and Efficient Test-Time Training

E2-TTT converts a token-wise recurrence with momentum and decay into a closed-form chunk-parallel transition under frozen chunk-start gradients. The resulting kernels preserve temporal gradient weighting while propagating only chunk-end states and retaining parallel computation.

  • E2-TTT maps token-wise updates with coupled momentum and decay to a chunk-parallel update that reproduces the chunk-end states without approximating per-token dynamics.
  • 3.1 Primal Formulation: Token-wise Dynamics: The target recurrence updates momentum from the previous momentum and current negative gradient, then applies momentum and weight decay to update fast weights.
  • 3.1 Primal Formulation: Token-wise Dynamics: Gradients are computed using fixed chunk-start weights, allowing token gradients to be computed in parallel before recurrent state updates.
  • Only chunk-end fast-weight and momentum states are propagated, while intermediate token states are not stored.
  • Naive sequential execution or full prefix scans can require O(C · d2) memory because dense states would be stored at every token.
  • The closed-form derivation compresses temporal dynamics into scalar coefficients that isolate each token’s contribution to the final state.
  • Momentum and weight kernels apply distinct temporal weighting, whereas LaCT uses only ηt for gradients and averages momentum across the chunk.
  • One backward pass reuses per-token activation gradients for both aggregates, adding only a reweighted aggregation over LaCT.

4 E2-TTT Model Architecture

E2-TTT combines sliding-window attention for local dependencies with a dynamically gated TTT branch for long-range context. Its input-dependent update coefficients and parallel branch fusion address the within-chunk blind spot of fixed-state output processing.

  • Motivation: Fixed-state chunk output prevents the TTT path from using immediate causal history within the current chunk.The hybrid design adds sliding-window attention to capture these local dependencies.
  • Overall Architecture: The architecture pairs E2-TTT with sliding-window attention to capture long-range and local dependencies, respectively.The two branches process shared queries, keys, and values before gated fusion.
  • Input-Dependent Dynamics: Input-dependent η_t, β_t, and γ_t modulate update magnitude, gradient persistence, and forgetting at each token.The coefficients are learned from the input through projections trained with the language-modeling loss.
  • Token Mixing Block: SWA and E2-TTT run in parallel over shared Q, K, V and are fused through a data-dependent gate.The fused output is then normalized and passed through the output projection.
  • Fast-weight Network: The fast-weight module is instantiated with either a GELU MLP or a SwiGLU MLP, both using residual connections and layer normalization.The GELU variant uses two weight matrices, W1 and W2.

W2 GELU(W1x)

The SwiGLU fast-weight variant uses a gated mechanism parameterized by three matrices.

  • The SwiGLU variant employs a gated mechanism parameterized by three matrices, {W1, W2, W3}.

5 Experiments

Experiments evaluate E2-TTT across language modeling, retrieval, length extrapolation, controlled update-rule comparisons, and video understanding. Results are strongest in retrieval and extrapolation, while language-modeling performance remains competitive.

  • Evaluation Setup: E2-TTT is evaluated on general language modeling, in-context retrieval, length extrapolation, controlled update-rule ablations, and long-context video understanding.Language models are trained at 340M and 1.3B parameters on 15B FineWeb-Edu tokens.
  • General Language Modeling: E2-TTT attains the lowest perplexity at both model scales, while average zero-shot accuracy differs by about two points across methods.At 1.3B parameters, E2-TTT_MLP reaches 54.5% average zero-shot accuracy versus 53.8% for the strongest baseline.
  • In-Context Retrieval: 43.6% average accuracy lets the 1.3B E2-TTT SwiGLU model surpass HQLT at 35.5% and LaCT at 36.7% on recall-intensive tasks.Full attention remains higher at 53.3% average, but E2-TTT narrows the gap.
  • Length Extrapolation: At 16K tokens, both E2-TTT variants retain at least 85% accuracy on S-NIAH-1, while E2-TTT_SwiGLU reaches 40.6% at 8K on S-NIAH-2.LaCT collapses to near-zero on S-NIAH-1 at 16K, and HQLT degrades to 25%.
  • Real-world Long Context Understanding: 14.1% average LongBench accuracy makes E2-TTT_SwiGLU the top model, outperforming Mamba2 on 13 of 14 tasks and LaCT on all 14.The MLP variant also exceeds all baselines with a 13.3% average.
  • Controlled Comparison: The controlled comparison isolates inner-loop aggregation by keeping the surrounding architecture fixed across the final rows of Table 5.The training context is 2K tokens, with columns beyond 2K reporting extrapolation.
  • Video Understanding: Qwen3VL-2B-Instruct plus E2-TTT improves VideoMMMU from 41.9 to 43.9 and LongVideoBench from 56.9 to 59.0 using only trained TTT parameters.These scores match or exceed the fully fine-tuned baseline trained on the same data.

6 Related Work

Test-Time Training treats recurrent state as learnable weights updated online to memorize in-context key-value associations. E2-TTT targets the sequential bottleneck by preserving per-token update structure while enabling chunk-level parallelism.

  • Test-Time Training: Test-Time Training represents recurrent state as learnable weights of an online-adapted nonlinear network that updates during inference.The fast weights use self-supervised losses to memorize in-context key-value associations.
  • Efficiency Trade-off: Token-wise TTT suffers from sequential hardware bottlenecks, whereas LaCT improves throughput by simplifying intra-chunk update aggregation.E2-TTT instead preserves temporal update structure while retaining chunk-level computational efficiency.

7 Conclusion

E2-TTT reconciles chunk-wise hardware efficiency with per-token expressivity by deriving a closed-form parallel transition for TTT states. The formulation preserves temporal update structure while enabling direct chunk-end computation.

  • Conclusion: E2-TTT reconciles chunk-wise processing efficiency with the per-token expressivity of TTT.Its closed-form scalar kernel parallelizes recurrent dynamics without the chunk-level approximations of prior chunk-wise methods.
  • Conclusion: The chunk-end momentum and weight states can be computed by unrolling and re-indexing the per-token recurrence into closed-form gradient, decay, and momentum contributions.The construction jumps from the chunk start to the chunk end without materializing intermediate states.
  • Conclusion: Cumulative products and scalar terms can be pre-computed in parallel, with log-space computation used for numerical stability at long chunk lengths.The final summation converts the cumulative products back to linear space.
  • Conclusion: The momentum state propagates each gradient to the chunk end through a suffix product, while the weight state combines decay, momentum, and gradient contributions.The momentum kernel has one decay channel and therefore does not require the cumulative ratio used by the weight kernel.

A.3 Stability of the Cumulative Kernel Rt

The cumulative kernel remains controlled across practical retention and decay regimes. Coupled decay and learning-rate parameterization also yields close agreement with the fully sequential reference.

  • A.3 Stability of the Cumulative Kernel Rt: 5.12 is the worst-case cumulative-kernel bound for ηbase = 10^-2 and C = 512 in the near-unit-retention regime.The bound follows from a per-token contribution bounded by ηbase · C.
  • A.3 Stability of the Cumulative Kernel Rt: When decay is aggressive, the cumulative ratio kernel is dominated by its first few terms and remains O(1).The suffix product for decay decreases geometrically when γt is small.
  • A.3 Stability of the Cumulative Kernel Rt: Coupling γt = 1−ηtαt makes larger updates increase forgetting, preventing carry inflation across chunks.This parameterization links update magnitude and retention rather than treating them as independent.
  • A.3 Stability of the Cumulative Kernel Rt: Relative ℓ2 deviation from the fully sequential reference stays below 2 × 10^-6 for both WC and MC over 128 chunks and 65,536 tokens.The reported verification uses C = 512.
  • A.3 Stability of the Cumulative Kernel Rt: A shared backward pass produces activation gradients for both state aggregates, requiring only one additional scalar-weighted contraction for MC.This keeps the extra momentum computation lightweight and is consistent with throughput close to LaCT.
  • A.3 Stability of the Cumulative Kernel Rt: Under frozen-W0 gradients, the chunk-end states are equivalently a token recurrence or scalar-weighted aggregation of per-token gradients.The optimization-style view emphasizes carried momentum across chunks, while the kernel view emphasizes within-chunk weighting.

B.2 Evaluation Details

The evaluation covers commonsense reasoning, in-context retrieval, length extrapolation, and video understanding, using established benchmarks and competitive hybrid baselines.

  • Commonsense Reasoning: Commonsense reasoning is evaluated on PIQA, HellaSwag, WinoGrande, ARC-e, ARC-c, Wikitext, and LAMBADA.Evaluations use lm-evaluation-harness.
  • In-context Retrieval: In-context retrieval is tested on FDA, SWDE, and SQuAD.SQuAD and SWDE use lm-evaluation-harness, while FDA follows the specified prior-work evaluation scripts.
  • Length Extrapolation: Length extrapolation uses PG19, GovReport, QMSum, NarrativeQA, Qasper, CodeParrot, and RULER needle-in-a-haystack tasks.The language-modeling context is 2K tokens, with S-NIAH-1 and S-NIAH-2 assessing effective capacity.
  • Video Understanding: Video understanding is evaluated on VideoMMMU and LongVideoBench with uniform frame sampling and task-specific frame and token budgets.VideoMMMU uses 512 frames and 256 tokens per frame; LongVideoBench uses 256 frames and 512 tokens.
  • Hybrid Baselines: Hybrid comparisons include HQLT and LaCT, using their specified synchronous or Muon-optimized variants.HQLT combines DeltaNet with window attention, while LaCT combines chunk-wise TTT with window attention.

B.3 Ablation Study

Ablations attribute gains to the closed-form kernel, hybrid branches, normalization, momentum, decay, and stable hyperparameter choices. The reported results also identify practical quality–efficiency trade-offs.

  • Effect of the closed-form kernel and chunk granularity: Replacing averaged chunk factors with the closed-form kernel lowers Wiki perplexity from 26.9 to 25.5 and raises average accuracy from 48.6% to 49.1% at C=512.The averaged baseline collapses per-token decay and momentum into chunk-level scalars.
  • Effect of the closed-form kernel and chunk granularity: C=512 further improves both metrics relative to C=1024, while smaller chunks narrow the gap to token-wise updates at a hardware-utilization cost.C=1024 already outperforms the averaged baseline, with Wiki ppl 25.6 and Avg 48.4%.
  • Branch contributions: 95.4% S-NIAH-1 accuracy at 8K and 93.6% at 16K require both sliding-window attention and the TTT branch.Sliding-window attention alone reaches 6.8% at 8K, while the TTT branch alone reaches 0.0% across the tested lengths.
  • Branch contributions: Post-chunk normalization trades S-NIAH-1 accuracy of 99.8 for 95.4 at 8K while improving LongBench from 13.4 to 14.1.Removing the fusion gate also reduces accuracy beyond 2K on both evaluation axes.
  • Inner-loop momentum and weight decay: Adding input-dependent momentum and weight decay reduces perplexity on both Wiki and LMB compared with the learning-rate-only baseline.The ablation uses C=512.
  • Sensitivity to base hyperparameters: Performance remains stable across the tested αbase range and for ηbase ∈ {10^-3, 10^-2}, while ηbase = 10^-1 degrades training.The defaults are ηbase = 10^-2 and αbase = 10^-1.

B.4 Analysis of the Learned Coefficients

Post-hoc measurements show that E2-TTT’s learned per-token coefficients vary substantially, align with their inputs, and contribute non-redundantly to retrieval and language-modeling behavior. The broader evaluation also reports efficient inference and strong extrapolation, while noting limits from the fixed within-chunk output state and tested scale.

  • Coefficient variation: The learned step size, forgetting rate, and decay rate are strongly input-dependent rather than near-constant.Median coefficients of variation are 0.95 for η_t, 0.46 for −log β_t, and 1.29 for 1 − γ_t; fewer than 1% of pairs fall below 0.1.
  • Coefficient variation: Over a 512-token chunk, the fast-weight carry β_t spans a factor of 129 between its 5th and 95th percentiles.This spread quantifies the temporal variation discarded by chunk-averaged updates.
  • Causal alignment: Shuffling coefficients across tokens hurts every intervention, with η alone raising perplexity by 0.45 and reducing S-NIAH-1 accuracy from 93.6% to 6.0% at 16K.The permutation preserves each coefficient’s marginal distribution while removing alignment with the token that produced it.
  • Causal alignment: The joint intervention costs +0.60 perplexity, close to the sum of the individual costs, supporting separate contributions from the coefficient heads.The paper reports individual costs of +0.45, +0.11, and +0.05, while noting that additivity is readable on perplexity only.
  • Efficiency: At 1.3B parameters, E2-TTT variants decode slightly faster than LaCT at comparable memory, although E2-TTTSwiGLU uses 8.08 versus 7.10 GB at batch 8.E2-TTTMLP uses 6.86 versus 7.10 GB at batch 8.
  • Extrapolation: E2-TTT maintains strong long-context retrieval, with E2-TTTSwiGLU reaching 93.6% on S-NIAH-1 at 16K while all baselines remain below 26%.On S-NIAH-2, E2-TTTSwiGLU remains the leading method through 16K, while sub-quadratic baselines fall below 30% from 8K onward.
Loading 2608.21308v1…