Rethinking Expressivity and Efficiency in Test-Time Training

Rethinking Expressivity and Efficiency in Test-Time Training

Zeyun Zhong1,3 Frederik Diederichs3

Joya Chen2 Manuel Martin3 Juergen Gall4,5 Juergen Beyerer1,3

arXiv:2608.21308v1 [cs.LG] 21 Aug 2026

1

Karlsruhe Institute of Technology (KIT) 2 National University of Singapore 3 Fraunhofer IOSB 4 Lamarr Institute for Machine Learning and Artificial Intelligence 5 University of Bonn

Abstract 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 E2 -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 E2 -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× the training context length. Meanwhile, E2 -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

State-of-the-art large language models (LLMs) [44, 16, 2] demonstrate remarkable capabilities in reasoning, coding, and multimodal understanding. However, once trained, the weights of these models remain frozen, leaving them to rely solely on accumulated memory, e.g., key-value cache [48, 49], for processing long-context agentic tasks [43, 41, 34]. This fundamental limitation hinders their ability to continuously learn from long-horizon tasks or generalize across tasks in real-world scenarios. Two primary approaches have emerged to address this limitation. The first seeks to reduce the computational complexity of attention from quadratic to linear by introducing RNN-style or hybrid architectures [13, 65, 47, 32], thereby enabling the processing of longer sequences. However, as demonstrated in Tandon et al. [58], these techniques still remain less effective than full attention in long-context language modeling. The second approach, known as Test-Time Training [57, 8], extends the concept of recurrent states to small, online-adapted sub-networks (e.g., fast weights [54]), allowing models to continuously adapt their parameters during inference. Numerous recent studies [12, 58, 69] have explored its potential for long-context text and multimodal tasks. Despite their promise, existing TTT methods face a well-known dilemma between expressivity and efficiency. On the one hand, vanilla TTT [57, 8] adopts token-wise updates, which offer powerful Preprint.

expressiveness but result in slow training, as token-level dependencies enforce sequential recurrence. On the other hand, chunk-wise TTT [69] substantially improves efficiency by executing updates over larger chunks, yet it sacrifices expressivity due to oversimplifications of the update rule, effectively obscuring the temporal variance of token importance within a chunk. To achieve the best of both worlds in expressivity and efficiency, we propose E2 -TTT, a novel TTT approach that reconciles token-level expressiveness with chunk-level computational efficiency. The core innovation is a closed-form parallel scalar kernel that enables parallel execution of recurrent dynamics with per-token learning rate, momentum, and decay. Specifically, our formulation preserves the exact per-token effect on the chunk-end fast-weight state, a temporal structure that prior chunkwise methods discard via averaging, while retaining chunk-level hardware parallelism. Furthermore, we integrate E2 -TTT into a hybrid architecture that combines TTT with sliding window attention, enabling our model to handle long-horizon tasks while maintaining fine-grained local context. We empirically validate E2 -TTT by training models of up to 1.3B parameters from scratch on 15B tokens from the HuggingFace FineWeb-Edu dataset [46], comparing against strong linear attention and hybrid baselines. Our evaluation spans general language modeling, in-context retrieval, and length extrapolation, where the gains concentrate in the retrieval and extrapolation regimes. Notably, our model exhibits strong robustness in length extrapolation, maintaining > 90% accuracy on passkey retrieval in the standard “Needle in a Haystack” test at 8× the training context, a regime where prior chunk-based methods collapse. Finally, we explore the versatility of our approach through parameterefficient adaptation for long-context video understanding, achieving performance comparable to full fine-tuning at the 2B-parameter scale on standard multimodal benchmarks. Overall, we demonstrate the benefit of expressive sequential updates while matching the training throughput of standard chunk-wise TTT, offering novel insights into the TTT paradigm.

2

Background: Test-Time Training

Unlike standard RNNs or linear attention models that compress history into an activation vector or matrix-valued states, Test-Time Training (TTT) [57] models the hidden state as learnable weights Wt of a sequence-dependent non-linear function fWt (·) : Rd → Rd . These weights, often referred to as fast weights [54], are rapidly adapted during both training and inference to dynamically store context, while the slow weights (i.e., model parameters) are fixed during inference. As illustrated in Figure 1, the general TTT process operates in two stages per step: an update stage and an output stage. First, input tokens are projected into query (qt ), key (kt ), and value (vt ) vectors, where qt , kt , vt ∈ Rd . The keys and values are used to update the fast weights via a selfsupervised objective, effectively encoding the context into a neural memory with fixed state size. Second, the updated fast weights are applied to the query to generate the output token ot . Based on how these update and output rules are structured, we categorize current TTT formulations into three groups, described below.

Figure 1: Conceptual overview of TTT. First, the key (kt ) and value (vt ) derived from the input are used to update (or “train”) the fast weights W, encoding the current context. Second, the updated weights Wt are applied to the query (qt ) to produce the output ot .

Expressive Token-wise Update. In this standard formulation, the update happens strictly at the token level. At each time step t, the fast weights Wt ∈ Rd×d are updated by minimizing a self-supervised reconstruction loss L between the transformed key fWt−1 (kt ) and the value vt . Let Gt denote the negative gradient, and ηt the learning rate:
Gt = −∇W L fWt−1 (kt ), vt , Wt = Wt−1 + ηt Gt . (1) The loss function is commonly Mean Squared Error [57]. The updated fast weights Wt are then immediately used to compute the output vector ot for the current query qt : ot = fWt (qt ). 2

(2)

Table 1: Comparison of TTT variants: vanilla TTT [57], Mini-Batch TTT [57, 8], and ChunkTTT [69]. E2 -TTT combines per-token update dynamics with the parallel training efficiency of chunk-based methods. † Gradients evaluated at the mini-batch/chunk-start weights. Update Forward Efficient Method Granularity Pass Training

TTT Mini-Batch TTT Chunk-TTT E2 -TTT (Ours)

Token-wise Token-wise† Chunk-wise Token-wise†

Token-wise Token-wise Chunk-wise Chunk-wise

✗ ✗ ✓ ✓

Token-wise Update with Mini-Batch Gradient. While the token-wise update offers high expressivity [57], it enforces sequential recurrence that prevents parallelization and results in memory-bound operations with poor hardware utilization. Consequently, recent work [57, 8] adapts a mini-batch approach. The sequence is partitioned into mini-batches indexed by r, each of small size B (e.g., 16). [r] We use the notation (·)t to denote a variable associated with the t-th token within the r-th batch. Let [r−1] [r] WB = W0 denote the state from the previous batch. Inside batch r, gradients are approximated [r−1] using this fixed WB , which decouples forward pass dependencies: [r] [r]
[r] [r] [r] [r] [r] [r] Gt = −∇W L fW[r−1] (kt ), vt , Wt = Wt−1 + ηt Gt , ot = fW[r] (qt ) (3) t

B

[r] This approximation enables the parallel computation of gradients Gt for all tokens t ∈ {1, . . . , B}. [r] Once calculated, the fast weights Wt can be computed, e.g., via a parallel associative scan [55].

Efficient Chunk-wise Update. Despite enabling parallel gradient calculation, the small mini-batch size and the token-wise dependency in mini-batch update still result in low hardware utilization (often < 5% FLOPs efficiency [69]). To address this, LaCT [69] transitions to a large-chunk formulation. The sequence is partitioned into large chunks indexed by r, each of size C ≫ B (e.g., C = 512). Within the r-th chunk, LaCT leverages the linearity of differentiation to aggregate the optimization objective. Instead of materializing gradients at every token step, it computes the gradient of the accumulated weighted loss, producing a single update direction G[r] for the entire chunk. When momentum is integrated, LaCT adopts a simplified strategy: averaging the time-dependent momentum [r] factors βt to maintain a chunk-level momentum buffer M[r] : X  C [r] [r] [r]
[r] G = −∇W ηt L fW[r−1] (kt ), vt , (4) t=1

M[r] =



C 1 X

C t=1

[r]

βt



M[r−1] + G[r] ,

W[r] = W[r−1] + M[r] .

(5)

This formulation allows maintaining chunk-level states (i.e., G[r] , M[r] , W[r] ) instead of token-level states, reducing the number of materialized fast weight states from C to 1 per chunk. As the new weights W[r] are only available after processing the entire chunk, the output O[r] (comprising outputs for all tokens in the chunk) is computed using the weights from the previous chunk: O[r] = fW[r−1] (Q[r] ).

3

(6)

Expressive and Efficient Test-Time Training

Current TTT variants present a stark trade-off: token-wise updates offer high optimization fidelity but suffer from sequential bottlenecks, while chunk-wise updates achieve hardware efficiency by simplifying the update rule, effectively obscuring the temporal variance of token importance within a chunk. To resolve this dichotomy, we introduce E2 -TTT (Expressive and Efficient TTT). Our key contribution is a novel theoretical formulation that maps a token-wise recurrence with coupled momentum and decay into a closed-form, chunk-parallel update, reproducing its chunk-end states without approximating the per-token dynamics. This allows us to retain the precise per-token temporal 3

weighting of each gradient contribution while matching the hardware efficiency of coarse-grained chunk processing. Table 1 summarizes this advantage. For simplicity, we omit the chunk indices r [r] in this section (e.g. Wt = Wt ). We index steps within a chunk as t ∈ {1, . . . , C}, where t = 0 [r−1] [r−1] denotes the state inherited from the previous chunk (W0 = WC , M0 = MC ). 3.1

Primal Formulation: Token-wise Dynamics

We first define our target optimization behavior. Unlike simplified chunk averages (Eq. 5), we aim to enforce precise, token-level recurrence with coupled L2 decay and momentum, a configuration proven beneficial in recent work [69, 65, 8]. Let Wt , Mt , Gt ∈ Rd×d denote the fast weights, momentum, and negative gradient (update direction), respectively. Let γt , βt ∈ (0, 1) be the scalar decay and momentum factors, and ηt > 0 be the learning rate. Our targeted sequential recurrence is defined as: Mt = βt Mt−1 + ηt Gt ,

Wt = γt Wt−1 + Mt .

(7)

The specific parameterizations of γt , βt , and ηt are predicted from the input xt independently of (W, M), as described in Section 4. As in the mini-batch setting (Eq. 3), the gradients Gt are [r−1] computed using the fixed weights from the previous chunk (WC ), which decouples them across t and allows them to be computed in parallel before the update step. The states actually propagated to [r] [r] [r] [r] the next chunk are the chunk-end states (WC , MC ); all intermediate (Wt , Mt ) for t < C are not stored. 3.2

Closed-Form Parallelization via Scalar Kernels

A naive execution of Eq. (7) requires sequential iteration (t = 1 → C). While parallel prefix scans could theoretically compute the states, they typically necessitate computing and storing Wt , Mt , and Gt at every time step. Since these are dense matrices, storing the full sequence incurs a prohibitive memory cost of O(C · d2 ), causing bottlenecks on modern accelerators, especially given the large chunk size C. To enable efficient training without simplifying the dynamics of Eq. (7), we derive a closed-form equivalent that “jumps” from the chunk start (t = 0) to the end (t = C) in a single step. Our insight is that while the matrix states (W, M) are expensive, the decay and momentum factors (γt , βt ) are scalars. By unrolling the recurrence and re-indexing the summation, we can compress the entire history of temporal dynamics into scalar coefficients that isolate the contribution of each gradient Gt to the final state (see Appendix A for the full derivation): WC =

C Y

j C  Y t C C C  Y C    X  X Y X Y γt W0 + βi M0 γq + ηt Gt βi γq . (8)

t=1

| {z } Decay term

t=1

q=t+1

i=1

{z

|

t=1

}

Momentum term

|

j=t

i=t+1

q=j+1

{z

Gradient term

}

The inner summation in the gradient term encapsulates the cumulative impact of the t-th token on the final chunk weight. We formalize this by defining efficient suffix products βet , γ et and a ratio sum Rt : βet :=

C Y

βi ,

γ et :=

i=t+1

C Y

γi ,

i=t+1

Rt :=

C X γ ei i=t βi

,

(9)

e

QC with the convention i=C+1 (·) = 1. These terms depend only on the scalars βt , γt and can therefore be computed efficiently in parallel (e.g., via log-space cumulative sums)
with negligible cost. Qj QC e Rewriting the prefix products in Eq. (8) via the identity i=1 βi = i=1 βi /βj and unrolling the momentum recurrence in parallel yields the following result. Proposition 3.1 (Exact chunk-end states). Under the recurrence of Eq. (7) with Gt evaluated at the frozen chunk-start W0 and scalars ηt , βt , γt independent of (W, M), the chunk-end states satisfy MC = βe0 M0 +

C X

KtM Gt ,

KtM := ηt βet ,

(10)

KtW := ηt βet Rt .

(11)

t=1

WC = γ e0 W0 + βe0 R1 M0 +

C X

KtW Gt ,

t=1

4

Algorithm 1 E2 -TTT chunk update. C C Require: chunk-start state (W0 , M0 ); tokens
{(qt , kt , vt )}t=1 ; scalars {ηt , βt , γt }t=1 1: ot = fW0 (qt ), Lt = L fW0 (kt ), vt // Compute outputs and per-token losses et , γ 2: Compute scalar aggregates β et , Rt via log-space cumulative sums; form KtW , KtM 3: gt ← −∂Lt /∂fW0 (kt ) // Per-token activation gradients (single backward)
P 4: ∆W ← AGGREGATE {KtW }, {gt }, {kt } // ≡ − t KtW ∇W Lt W 0
P 5: ∆M ← AGGREGATE {KtM }, {gt }, {kt } // ≡ − t KtM ∇W Lt W ; reuses {gt } from line 3 0 6: WC ← γ e0 W0 + βe0 R1 M0 + ∆W e0 M0 + ∆M 7: MC ← β 8: return outputs {ot }, next-chunk state (WC , MC )

Two scalar kernels with distinct temporal weighting. Proposition 3.1 reveals a key structural fact: the chunk-end weights WC and the chunk-end momentum MC aggregate the per-token gradients {Gt } through two distinct scalar kernels. The momentum kernel KtM = ηt βet measures how much of Gt survives to the chunk end through the momentum decay βet . The weight kernel KtW = ηt βet Rt additionally weights this by the cumulative weight-decay survival Rt , capturing how the gradient propagates through both the momentum buffer and the subsequent weight updates. A bound analysis of Rt and the resulting kernels KtW , KtM under near-unit βt , γt is provided in Appendix A.3. This is in stark contrast to LaCT [69], where the gradient term carries only ηt (without βet or Rt ) and the momentum factor is collapsed to a chunk average, eliminating the temporal structure of the aggregation (Eq. 5). P P One backward pass for both aggregates. Computing the two aggregates t KtW Gt and t KtM Gt would naively require two backward passes. We collapse them to one by noting that the kernels differ only in scalar coefficients and share the per-token activation gradients gt = −∂Lt /∂fW0 (kt ) (see Appendix A.4). These gradients arise as a C × d intermediate of LaCT’s backward, which contracts them to a d × d weight gradient. We retain them across both aggregations rather than discard them after the first. The marginal cost over LaCT is therefore a single re-weighted aggregation. Algorithm 1 summarizes the procedure as a direct instantiation of Proposition 3.1. The output is computed using the previous chunk’s weights, O[r] = fW[r−1] (Q[r] ), consistent with chunk-wise TTT and ensuring a fully parallel forward pass. Algorithm 1 retains LaCT’s O(C · d2 ) asymptotic time complexity per chunk, which we confirm empirically in Appendix B.5. Cross-chunk equivalence to the original sequential recurrence follows by induction, since (WC , MC ) is precisely the state propagated to the next chunk; we numerically verify this on a 128-chunk trajectory (C = 512, 65,536 tokens), where the relative L2 deviation from the fully sequential reference is below 2 × 10−6 on both WC and MC , within FP32 round-off accumulated over the trajectory.

4

E2 -TTT Model Architecture

The design of our architecture is driven by a specific characteristic of chunk-wise processing: the output step (Eq. 6) applies a fixed state to the entire chunk. This renders the TTT path “blind” to the [r] immediate causal history within the current chunk (i.e., token xt cannot influence the processing [r] of xt+1 via the fast weights). To address this, we adopt the proven hybrid design pattern in prior work [4, 26, 42, 69], pairing our E2 -TTT module (for modeling non-local dependencies) with a standard Sliding Window Attention module (for capturing local dependencies), as shown in Fig. 2. Overall Architecture. The macro-architecture of E2 -TTT follows the standard Llama design [59], interleaving token mixing layers with SwiGLU MLP blocks. Given an input sequence X ∈ RL×d , we first project it into shared queries Q, keys K, and values V. These are fed into two parallel processing branches: a Sliding Window Attention (SWA) module for local retrieval, and our E2 -TTT layer for long-range context modeling. For the TTT path, we follow [28, 69] by applying SiLU activation to the queries and keys, followed by L2 normalization to ensure training stability. Simultaneously, the scalars required for the TTT update rule are predicted by the network directly: the learning rate ηt , momentum βt , and decay strength αt are projected directly from X via lightweight linear heads, 5

while the final decay factor γt is derived from ηt and αt . To combine the complementary strengths of both branches, we adopt a dynamic mixing strategy [28]: the outputs are fused via element-wise interpolation, controlled by a data-dependent gate projected from the input X. Finally, the fused output is normalized and processed through the output projection. Input-Dependent Dynamics. To allow the model to modulate its learning dynamics at the token level, we parameterize the update coefficients ηt , βt , γt as functions of the input xt ∈ Rd . We adopt standard parameterization schemes from prior TTT and linear attention literature to ensure stability and effective longterm retention. The corresponding linear projections are learned end-to-end from the languagemodeling loss alone, without auxiliary objectives. Specific base scalars are detailed in Appendix B.1. Performance is empirically robust to these base scalars across an order of magnitude (Appendix B.3). Appendix B.4 further analyzes the coefficients learned by these projections. The learning rate ηt modulates the update magnitude. We parameterize it as a sigmoid-gated Figure 2: E2 -TTT token mixing block. SWA and value scaled by a base learning rate ηbase [57]: E2 -TTT run in parallel over shared Q, K, V and ηt = ηbase · σ(Linearη (xt )). (12) are fused by a data-dependent gate. The TTT upThe momentum factor βt determines the persis- date is modulated by input-dependent α, η, β. tence of historical gradients. To encourage long-term memory retention, we bias the momentum towards values close to 1 using a timescale factor τ [64, 56]: 1

βt = σ(Linearβ (xt )) τ .

(13)

The decay factor γt serves as a forget gate. To prevent excessive regularization when the update step is small, we derive γt from a dynamic decay strength αt (scaled by αbase ) coupled with the learning rate ηt , consistent with standard optimization practices [39]: αt = αbase · σ(Linearα (xt )),

γt = 1 − ηt αt .

(14)

Fast-weight Network. While our framework supports any neural network as the fast-weight module f , we instantiate f using two standard MLP configurations [57, 69]: a standard GELU MLP and a SwiGLU MLP. Both variants incorporate residual connections and Layer Normalization (LN). The GELU MLP variant is parameterized by two weight matrices W1 and W2 . The forward transformation for an input x is defined as:
fW (x) = x + LN W2 GELU(W1 x) . (15) The SwiGLU variant employs a gated mechanism parameterized by three matrices {W1 , W2 , W3 }:
fW (x) = x + LN W3 (SiLU(W1 x) ⊙ (W2 x)) . (16)

5

Experiments

We evaluate our method across three distinct categories of tasks to verify its versatility and robustness: general language modeling, to validate overall capabilities on standard benchmarks (Sec. 5.1); incontext retrieval, to assess the model’s precise recall abilities (Sec. 5.2); and length extrapolation, to test generalization to sequence lengths exceeding the training horizon (Sec. 5.3). We then isolate the contribution of the update rule itself in a tightly matched ablation (Sec. 5.4). Finally, we extend our evaluation to long-context video understanding, exploring the potential of our E2 -TTT to leverage long-range temporal dependencies in multimodal data (Sec. 5.5). Setup. Following HQLT [28], we train language models with either 340M or 1.3B trainable parameters (using sequence lengths of 2048 and 2240, respectively) from scratch on 15B tokens of the HuggingFace FineWeb-Edu dataset [46]. These choices mostly follow the configurations of the recent work on linear transformers [64, 66, 20], including the choice of tokenizer 6

Table 2: General language modeling experiments. We report perplexity (↓) and zero-shot accuracy (↑) for models with 340M and 1.3B parameters, all trained on 15B tokens. E2 -TTT attains the lowest perplexity at both scales; on average zero-shot accuracy all methods here lie within about two points of one another. Results marked with * are cited from HQLT. Model

Attn. Wiki. LMB. LMB. PIQA Hella. Wino. ARC-e ARC-c Avg. window ppl ↓ ppl ↓ acc ↑ acc ↑ acc_n ↑ acc ↑ acc ↑ acc_n ↑ ↑

340M params Transformer++* DeltaNet* HQLT* LaCT E2 -TTTSwiGLU E2 -TTTMLP

2048 – 512 512 512 512

26.5 27.6 27.0 27.0 26.4 25.5

34.9 35.0 28.0 28.2 27.8 25.0

33.9 32.8 35.9 37.1 36.6 37.0

67.6 67.1 66.3 66.8 67.7 67.4

41.0 40.8 41.3 41.6 42.2 43.9

53.7 52.6 53.2 51.4 51.6 54.9

60.2 58.5 60.1 59.0 61.4 61.3

29.0 28.8 29.0 30.0 28.3 30.2

47.6 46.8 47.6 47.7 48.0 49.1

1.3B params Transformer++* DeltaNet* Mamba2 HQLT LaCT E2 -TTTSwiGLU E2 -TTTMLP

2048 – – 512 512 512 512

19.8 20.6 20.5 20.3 20.1 19.9 19.5

17.9 19.9 17.0 16.1 16.7 15.8 15.3

42.6 39.3 41.1 42.6 42.0 43.2 43.7

71.0 70.1 71.0 71.0 70.4 71.1 71.4

50.3 49.5 51.5 51.1 51.6 51.9 52.7

55.8 52.5 54.5 55.1 55.4 54.2 56.1

65.2 68.5 66.5 69.2 65.6 67.5 67.6

33.2 34.2 35.0 34.0 33.1 34.1 35.2

53.0 52.3 53.3 53.8 53.0 53.6 54.5

(fla-hub/transformer-1.3B-100B). The subword-unit vocabulary size is 32K for all models. All models have 24 layers. The baseline Transformer architecture is from Touvron et al. [59] (denoted as Transformer++, following prior convention) and the DeltaNet configuration is from Yang et al. [66]. We also compare against two hybrid models: HQLT and LaCT [69], both of which utilize window attention combined with DeltaNet and chunk-wise TTT, respectively. For all baseline models, we either report the results from the respective papers or train them using official public repositories under identical conditions to E2 -TTT for fair comparison. We use a chunk size of 512 for both LaCT and our models. Full architectural, training, and evaluation details are in Appendix B.1–B.2; a separate feasibility run at a larger budget is reported in Appendix B.7. 5.1

General Language Modeling Capabilities

In Table 2, we present the language modeling perplexity and zero-shot accuracy on commonsense reasoning benchmarks for models with 340M and 1.3B parameters. E2 -TTT obtains the lowest perplexity at both scales, with the margin most visible on LAMBADA (15.3 vs. 16.1 for the strongest baseline at 1.3B, and 25.0 vs. 28.0 at 340M). On zero-shot accuracy the picture is more modest: E2 -TTTMLP reaches 54.5% average at 1.3B against 53.8% for the strongest baseline, and leads on five of six tasks while trailing HQLT and DeltaNet on ARC-e. 5.2

Evaluating In-Context Retrieval Abilities

In this section, we evaluate the performance Table 3: Accuracy on real-world retrieval tasks with of E2 -TTT on recall-intensive tasks. Follow- input truncated to 2K tokens. ∗ : from HQLT. ing prior work [66, 3, 28], we focus on three Model Attn. SWDE SQuAD FDA Avg. window acc ↑ acc ↑ acc ↑ tasks: FDA [3], SWDE [38], and SQuAD [51]. Table 3 summarizes the results for 340M params Transformer++* 2048 44.9 36.9 52.3 44.7 the 340M and 1.3B parameter scales. ConTransformer++* 1024 30.4 25.5 31.2 29.0 sistent with prior observations by Irie et al. DeltaNet* – 18.5 25.2 8.6 17.4 [28], full attention models (Transformer++) HQLT* 512 22.9 35.7 17.3 25.3 still hold an advantage on these tasks (53.3% LaCT 512 25.0 39.0 22.2 28.7 E2 -TTTSwiGLU 512 35.4 39.6 26.1 33.7 average) due to their explicit history access. E2 -TTTMLP 512 23.6 40.1 26.0 29.9 However, E2 -TTT significantly narrows this gap compared to prior works, for instance, 1.3B params Transformer++* 2048 53.7 41.5 64.7 53.3 our 1.3B SwiGLU model (43.6%) surpasses DeltaNet* – 32.9 29.9 23.6 28.8 HQLT (35.5%) and LaCT (36.7%) on average, Mamba2 – 33.8 37.3 26.1 32.4 demonstrating the benefit of our expressive HQLT 512 32.1 43.4 31.0 35.5 LaCT 512 39.2 44.9 26.0 36.7 chunk-wise update rule. Furthermore, we obE2 -TTTSwiGLU 512 47.3 44.4 39.0 43.6 serve that our SwiGLU variant outperforms E2 -TTTMLP 512 37.2 45.6 28.5 37.1 the MLP variant by a large margin (43.6% vs 7

GovReport

Loss

Extrapolation

4.5

2.5

3

2.0

2

4.0 3.5 3.0

Qasper

5

Loss

NarrativeQA

5.0

Extrapolation

4

3.0

2.5

CodeParrot

4.0

Extrapolation

4 3 2

QMSum

5

Extrapolation

3.5

3.5

4.0

3.0

3.5

2.5 0

2k

4k

6k

8k

2.0

PG19

4.5

Extrapolation

Extrapolation

3.0 0

2k

Length

4k

6k

8k

2.5

0

Length HQLT

LaCT

Ours SwiGLU

2k

4k

6k

8k

Length Ours MLP

Figure 3: Length extrapolation of 1.3B models. We visualize the per-token loss metric [63] on six long-context benchmarks, where E2 -TTT maintains stability while the LaCT baseline diverges. 37.1% average), validating our hypothesis that complex gating mechanisms are beneficial for precise information extraction. 5.3

Evaluating Length Extrapolation Capabilities

A critical advantage of our E2 -TTT is its potential to generalize to sequence lengths significantly exceeding the training horizon. We evaluate this capability across three dimensions: language modeling stability, synthetic long-context retrieval, and real-world long-context benchmarks. Language Modeling Extrapolation. We first assess whether the models can maintain stable language modeling performance when processing sequences longer than their training context (2K tokens). Figure 3 visualizes the per-token loss metric [63] on six long-context datasets. As shown in Fig. 3, LaCT struggles with length extrapolation, with loss exploding immediately past the training boundary on all six datasets. HQLT maintains stability but shows higher validation loss (e.g., on CodeParrot). In contrast, both E2 -TTT variants demonstrate lower and stable loss curves throughout the extrapolated region. This confirms that our expressive chunk-wise updates robustly generalize to unseen lengths.

Accuracy [%]

S-NIAH-1 S-NIAH-2 Long-range Associative Recall. Sta100 100 Extrapolation bility alone does not guarantee in75 75 formation retention. To test effec50 50 tive capacity, we employ the “Single25 25 Needle in a Haystack” (S-NIAH) Extrapolation benchmark [24], extending the test se0 0 512 1k 2k 4k 8k 16k 512 1k 2k 4k 8k 16k quence up to 16K tokens (8× the trainLength Length ing horizon). Figure 4 reveals a stark HQLT LaCT Ours SwiGLU Ours MLP contrast in extrapolation behavior. On Figure 4: Zero-shot accuracy on S-NIAH-1 and S-NIAH-2 S-NIAH-1, both E2 -TTT variants refor 1.3B models. All methods use sliding window attention tain at least 85% accuracy at 16K, with a 512-token window. E2 -TTT maintains robust retrieval while LaCT collapses to near-zero and well beyond the training context, while LaCT collapses past HQLT degrades to 25%. On S-NIAH4K and HQLT degrades steadily across both tasks. 2 2, E -TTTSwiGLU retains a clear lead through 16K, achieving 40.6% at 8K versus 15.4% (HQLT) and 5.6% (LaCT). The SwiGLU variant consistently outperforms the MLP variant in the extrapolation regime, validating the benefit of gated updates for handling complex retrieval patterns. A more detailed numerical comparison, including additional baselines, is provided in Appendix B.6.

Real-world Long Context Understanding. Finally, we evaluate performance on LongBench [7], a comprehensive suite of 14 real-world long-context tasks. As detailed in Table 4, our method translates its length extrapolation advantages into practical gains. E2 -TTTSwiGLU achieves the highest average score (14.1%), outperforming HQLT (12.1%) and the SSM-based Mamba2 (10.3%), while nearly doubling the performance of LaCT (7.7%). The improvements are systematic across task categories: E2 -TTTSwiGLU outperforms Mamba2 on 13 of 14 tasks, HQLT on 11 of 14, and LaCT on all 14. The MLP variant similarly outperforms all baselines (13.3% average), confirming that the gains stem from our update rule rather than a specific gating choice. This consistency across diverse task types—single8

Table 4: Evaluation on Longbench. Accuracy on 14 tasks from LongBench: Narrative QA, QasperQA, MultiField QA, HotpotQA, 2WikiMulti QA, Musique, GovReport, QMSum, MultiNews, TRec, Trivia QA, SamSum, LCC, and RepoBench-P by order. Model (1.3B)

Single-Doc QA Multi-Doc QA Summarization Few-shot Code Avg. NQA QQA MFQ HQA 2WM Mus GvR QMS MNs TRC TQA SSM LCC RBP

Mamba2 HQLT LaCT E2 -TTTSwiGLU E2 -TTTMLP

2.0 0.8 1.1 2.8 1.7

4.6 1.0 4.5 4.9 4.4

11.4 9.3 8.2 14.2 12.4

3.5 3.3 3.5 5.0 4.4

8.9 7.9 4.7 9.0 10.4

2.8 2.0 2.4 3.8 3.1

4.7 6.7 9.4 10.3 8.8

16.6 14.5 5.1 16.7 16.6

7.7 12.4 11.5 12.1 13.6

9.5 25.8 10.0 34.0 26.5

16.6 19.9 5.6 16.2 20.3

13.0 23.1 7.6 20.1 19.3

22.3 19.8 19.4 24.4 19.5

20.0 22.8 14.5 23.2 24.6

10.3 12.1 7.7 14.1 13.3

document QA, multi-document QA, summarization, few-shot, and code—demonstrates that our expressive chunk-wise update rule enables E2 -TTT to leverage extended context more effectively than baselines spanning multiple sub-quadratic paradigms (SSM, linear attention, and TTT).

5.4

Ablation Study

To isolate the update rule from the sur- Table 5: Controlled comparison of update rules (1.3B). rounding architecture, we train two addi- The last three rows share our framework and differ only tional 1.3B models inside our own frame- in how the inner loop is aggregated. Training context is work. Apart from the published LaCT con- 2K tokens. Columns to the right of 2K are extrapolation. figuration and the SWA-only reference, all S-NIAH-1 S-NIAH-2 rows of Table 5 share the same architecture, Method 2K 4K 8K 16K 2K 4K 8K including SWA and the fusion gate, and differ only in how the inner loop is aggregated. SWA only 26.8 – 6.8 2.8 – – – LaCT-Matched ports LaCT’s update rule LaCT 100.0 81.0 37.4 3.0 91.4 80.6 5.6 into our framework. Chunk-averaged is our LaCT-Matched 99.8 99.2 51.6 0.0 98.6 65.8 4.8 Chunk-averaged 25.2 10.8 6.8 3.0 30.4 18.2 7.2 own model with the per-token momentum E2 -TTTSwiGLU 95.6 95.0 95.4 93.6 96.8 87.6 40.6 and decay factors collapsed to chunk-level scalars. LaCT-Matched improves over its default configuration on S-NIAH-1 within and near the training context (81.0→99.2 at 4K) and is slightly ahead of our own model at 2K, while on S-NIAH-2 it is mixed (98.6 vs. 91.4 at 2K, 65.8 vs. 80.6 at 4K). Beyond the training context, however, it still collapses to 0.0 at 16K. Chunk-averaged isolates the design choice itself: collapsing the per-token scalars drops S-NIAH-1 at 8K from 95.4 to 6.8, the level of the SWA baseline alone, and the degradation already begins inside the training context (95.6→25.2 at 2K). This is not a training failure: in-domain, the collapsed model is marginally ahead, at 15.6 vs. 15.8 LAMBADA perplexity and 53.8% vs. 53.6% zero-shot average. We therefore attribute the long-context retrieval gains, both within and beyond the training context, to our expressive chunk-wise update rule. Ablations of the remaining components, including inner-loop optimizer terms, chunk granularity, and sensitivity to the base scalars, are reported in Appendix B.3.

5.5

Multimodal Evaluation

To assess generalization beyond lan- Table 6: Video understanding benchmarks. Integrating E2 guage modeling, we integrate E2 - TTT into Qwen3VL-2B-Instruct and training only the TTT TTT into Qwen3VL-2B-Instruct [6] parameters matches full fine-tuning of the backbone on the as a parallel branch fused with self- same training data. attention layers. We freeze the preVideoMMMU LongVideoBench trained backbone and train only the Model Perc. Comp. Adap. Avg. Perc. Rel. Avg. TTT parameters on a 43K-sample subset of LLaVA-Video-178K [70]. Qwen3VL – – – 41.9 63.1 51.5 56.9 Qwen3VL (SFT) 60.3 37.3 33.0 43.5 63.4 54.0 58.3 Qwen3VL+E2 -TTT improves over Qwen3VL+E2 -TTT 60.7 37.3 33.7 43.9 63.7 54.9 59.0 the frozen base model (43.9 vs. 41.9 on VideoMMMU, 59.0 vs. 56.9 on LongVideoBench) and matches a fully fine-tuned baseline trained on the same data (43.5 and 58.3, respectively). Additional training and evaluation details are provided in Appendix B.1 and B.2. 9

6

Related Work

Test-Time Training (TTT) [57] is an emerging paradigm in sequence modeling that redefines the recurrent states in RNNs as learnable weights of an online-adapted non-linear neural network. These weights, often referred to as fast weights [54], are continuously updated to learn in-context. Existing methods typically employ a self-supervised loss that encourages these fast weights to memorize key-value associations from in-context tokens, effectively treating inference as an online learning problem using variants of gradient descent. While TTT [57, 61] has opened a vast design space for novel recurrent model architectures [8, 30, 9, 58], the resulting sequential bottleneck leads to poor hardware utilization, limiting their scalability. Recent efforts like LaCT [69] have attempted to mitigate this by shifting from token-wise to chunk-wise updates, sacrificing intra-chunk dependency modeling to improve throughput. However, these methods often resort to simplified aggregation rules (e.g., flat averaging) that degrade the optimization fidelity. Our work builds upon this foundation but distinguishes itself by introducing E2 -TTT (Expressive and Efficient TTT), a formulation that reconciles the hardware efficiency of chunking with the mathematical precision of complex, tokenwise update dynamics. Combining global recurrence with local attention is a proven strategy for efficient long-context modeling [36, 52, 15]. Following recent work [28, 69], we adopt a within-layer hybrid architecture that shares the key, value, and query vectors between local attention and global recurrence. This design pattern has been successfully deployed in various forms [26, 42, 3]. However, prior instantiations typically rely on standard linear attention [31, 54, 66, 28] or simplified TTT update rules [69]. Our work retains this established architectural template but replaces the recurrence module with our E2 -TTT, empirically showing superior results.

7

Conclusion

We presented E2 -TTT (Expressive and Efficient TTT), a method that reconciles the hardware efficiency of chunk-wise processing with the per-token expressivity of the update rule in Test-Time Training. By deriving a closed-form scalar kernel, we enabled the parallel execution of per-token recurrent dynamics without the chunk-level approximations of prior chunk-wise methods. Empirically, E2 -TTT is on par with strong sub-quadratic baselines in language modeling and outperforms them on retrieval while demonstrating superior length extrapolation, sustaining > 90% passkey retrieval accuracy on “Needle in a Haystack” at 8× the training context length, where baseline methods collapse. These results suggest that preserving precise temporal dynamics is essential for long-context generalization, positioning E2 -TTT as a promising direction for scalable Test-Time Training.

Acknowledgements This work was supported by the JuBot project funded by the Carl-Zeiss-Foundation. The authors acknowledge support by the state of Baden-Württemberg through bwHPC. We acknowledge EuroHPC Joint Undertaking for awarding us access to MareNostrum5 at BSC, Spain. Experiments were performed on the HoreKa supercomputer, funded by the Ministry of Science, Research and the Arts Baden-Württemberg and by the Federal Ministry of Education and Research, and on MareNostrum5. Juergen Gall has been supported by the ERC Consolidator Grant FORHUE (101044724).

References [1] Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, Iain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katherine Millican, Malcolm Reynolds, et al. Flamingo: a visual language model for few-shot learning. Advances in neural information processing systems, 35:23716–23736, 2022. 18 [2] Anthropic. System card: Claude opus 4.5, 2025. 1 [3] Simran Arora, Brandon Yang, Sabri Eyuboglu, Avanika Narayan, Andrew Hojel, Immanuel Trummer, and Christopher Ré. Language models enable simple systems for generating structured views of heterogeneous data lakes. arXiv preprint arXiv:2304.09433, 2023. 7, 10, 18 [4] Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall-throughput tradeoff. arXiv preprint arXiv:2402.18668, 2024. 5

10

[5] Simran Arora, Aman Timalsina, Aaryan Singhal, Benjamin Spector, Sabri Eyuboglu, Xinyi Zhao, Ashish Rao, Atri Rudra, and Christopher Ré. Just read twice: closing the recall gap for recurrent language models. arXiv preprint arXiv:2407.05483, 2024. 18 [6] Shuai Bai, Yuxuan Cai, Ruizhe Chen, Keqin Chen, Xionghui Chen, Zesen Cheng, et al. Qwen3-vl technical report. arXiv preprint arXiv:2511.21631, 2025. 9, 19 [7] Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd annual meeting of the association for computational linguistics (volume 1: Long papers), pages 3119–3137, 2024. 8, 18 [8] Ali Behrouz, Peilin Zhong, and Vahab Mirrokni. Titans: Learning to memorize at test time. arXiv preprint arXiv:2501.00663, 2024. 1, 3, 4, 10, 17, 21 [9] Ali Behrouz, Zeman Li, Praneeth Kacham, Majid Daliri, Yuan Deng, Peilin Zhong, Meisam Razaviyayn, and Vahab Mirrokni. Atlas: Learning to optimally memorize the context at test time. arXiv preprint arXiv:2505.23735, 2025. 10 [10] Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, pages 7432–7439, 2020. 18 [11] Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457, 2018. 18 [12] Karan Dalal, Daniel Koceja, Jiarui Xu, Yue Zhao, Shihao Han, Ka Chun Cheung, Jan Kautz, Yejin Choi, Yu Sun, and Xiaolong Wang. One-minute video generation with test-time training. In Proceedings of the Computer Vision and Pattern Recognition Conference, pages 17702–17711, 2025. 1, 18 [13] Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv:2405.21060, 2024. 1, 22 [14] Pradeep Dasigi, Kyle Lo, Iz Beltagy, Arman Cohan, Noah A Smith, and Matt Gardner. A dataset of information-seeking questions and answers anchored in research papers. arXiv preprint arXiv:2105.03011, 2021. 18 [15] Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mixing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402.19427, 2024. 10 [16] Google DeepMind. Gemini 3 pro model card, 2025. 1 [17] Alexander Richard Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir Radev. Multi-news: A large-scale multi-document summarization dataset and abstractive hierarchical model. In Proceedings of the 57th annual meeting of the association for computational linguistics, pages 1074–1084, 2019. 18 [18] Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, et al. A framework for few-shot language model evaluation. Zenodo, 2021. 18, 19 [19] Bogdan Gliwa, Iwona Mochol, Maciej Biesek, and Aleksander Wawer. Samsum corpus: A humanannotated dialogue dataset for abstractive summarization. arXiv preprint arXiv:1911.12237, 2019. 19 [20] Riccardo Grazzi, Julien Siems, Arber Zela, Jörg KH Franke, Frank Hutter, and Massimiliano Pontil. Unlocking state-tracking in linear rnns through negative eigenvalues. arXiv preprint arXiv:2411.12537, 2024. 6 [21] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First conference on language modeling, 2024. 18 [22] Daya Guo, Canwen Xu, Nan Duan, Jian Yin, and Julian McAuley. Longcoder: A long-range pre-trained language model for code completion. In International Conference on Machine Learning, pages 12098– 12107. PMLR, 2023. 19 [23] Xanh Ho, Anh-Khoa Duong Nguyen, Saku Sugawara, and Akiko Aizawa. Constructing a multi-hop qa dataset for comprehensive evaluation of reasoning steps. arXiv preprint arXiv:2011.01060, 2020. 18

11

[24] Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What’s the real context size of your long-context language models? arXiv preprint arXiv:2404.06654, 2024. 8, 18 [25] Kairui Hu, Penghao Wu, Fanyi Pu, Wang Xiao, Yuanhan Zhang, Xiang Yue, Bo Li, and Ziwei Liu. Video-mmmu: Evaluating knowledge acquisition from multi-discipline professional videos. arXiv preprint arXiv:2501.13826, 2025. 19 [26] Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In International conference on machine learning, pages 9099–9117. PMLR, 2022. 5, 10 [27] Luyang Huang, Shuyang Cao, Nikolaus Parulian, Heng Ji, and Lu Wang. Efficient attentions for long document summarization. arXiv preprint arXiv:2104.02112, 2021. 18 [28] Kazuki Irie, Morris Yau, and Samuel J Gershman. Blending complementary memory systems in hybrid quadratic-linear transformers. arXiv preprint arXiv:2506.00744, 2025. 5, 6, 7, 10, 18, 19, 23 [29] Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. arXiv preprint arXiv:1705.03551, 2017. 18 [30] Mahdi Karami, Razvan Pascanu, and Vahab Mirrokni. Lattice: Learning to efficiently compress the memory. arXiv preprint arXiv:2504.05646, 2025. 10 [31] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156–5165. PMLR, 2020. 10 [32] Kimi Team, Yu Zhang, Zongyu Lin, Xingcheng Yao, Jiaxi Hu, Fanqing Meng, Chengyin Liu, Xin Men, Songlin Yang, Zhiyuan Li, Wentao Li, Enzhe Lu, Weizhou Liu, Yanru Chen, Weixin Xu, Longhui Yu, Yejie Wang, Yu Fan, Longguang Zhong, Enming Yuan, Dehao Zhang, Yizhi Zhang, T. Y. Liu, Haiming Wang, Shengjun Fang, Weiran He, Shaowei Liu, Yiwei Li, Jianlin Su, Jiezhong Qiu, Bo Pang, Junjie Yan, Zhejun Jiang, Weixiao Huang, Bohong Yin, Jiacheng You, Chu Wei, Zhengtao Wang, Chao Hong, Yutian Chen, Guanduo Chen, Yucheng Wang, Huabin Zheng, Feng Wang, Yibo Liu, Mengnan Dong, Zheng Zhang, Siyuan Pan, Wenhao Wu, Yuhao Wu, Longyu Guan, Jiawen Tao, Guohong Fu, Xinran Xu, Yuzhi Wang, Guokun Lai, Yuxin Wu, Xinyu Zhou, Zhilin Yang, and Yulun Du. Kimi linear: An expressive, efficient attention architecture. arXiv:2510.26692, 2025. 1 [33] Tomáš Kočiskỳ, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. The narrativeqa reading comprehension challenge. Transactions of the Association for Computational Linguistics, 6:317–328, 2018. 18 [34] Junlong Li, Wenshuo Zhao, Jian Zhao, Weihao Zeng, Haoze Wu, Xiaochen Wang, Rui Ge, Yuxuan Cao, Yuzhen Huang, Wei Liu, Junteng Liu, Zhaochen Su, Yiyang Guo, Fan Zhou, Lueyang Zhang, Juan Michelini, Xingyao Wang, Xiang Yue, Shuyan Zhou, Graham Neubig, and Junxian He. The tool decathlon: Benchmarking language agents for diverse, realistic, and long-horizon task execution, 2025. 1 [35] Xin Li and Dan Roth. Learning question classifiers. In COLING 2002: The 19th International Conference on Computational Linguistics, 2002. 18 [36] Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, et al. Jamba: A hybrid transformer-mamba language model. arXiv preprint arXiv:2403.19887, 2024. 10 [37] Tianyang Liu, Canwen Xu, and Julian McAuley. Repobench: Benchmarking repository-level code auto-completion systems. arXiv preprint arXiv:2306.03091, 2023. 19 [38] Colin Lockard, Prashant Shiralkar, and Xin Luna Dong. Openceres: When open information extraction meets the semi-structured web. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 3047–3056, 2019. 7, 18 [39] Ilya Loshchilov and Frank Hutter. arXiv:1711.05101, 2017. 6

Decoupled weight decay regularization.

arXiv preprint

[40] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv:1609.07843, 2016. 18

12

[41] Mike A. Merrill, Alexander G. Shaw, Nicholas Carlini, Boxuan Li, Harsh Raj, Ivan Bercovich, Lin Shi, Jeong Yeon Shin, Thomas Walshe, E. Kelly Buchanan, Junhong Shen, Guanghao Ye, Haowei Lin, Jason Poulos, Maoyu Wang, Marianna Nezhurina, Jenia Jitsev, Di Lu, Orfeas Menis Mastromichalakis, Zhiwei Xu, Zizhao Chen, Yue Liu, Robert Zhang, Leon Liangyu Chen, Anurag Kashyap, Jan-Lucas Uslu, Jeffrey Li, Jianbo Wu, Minghao Yan, Song Bian, Vedang Sharma, Ke Sun, Steven Dillmann, Akshay Anand, Andrew Lanpouthakoun, Bardia Koopah, Changran Hu, Etash Guha, Gabriel H. S. Dreiman, Jiacheng Zhu, Karl Krauth, Li Zhong, Niklas Muennighoff, Robert Amanfu, Shangyin Tan, Shreyas Pimpalgaonkar, Tushar Aggarwal, Xiangning Lin, Xin Lan, Xuandong Zhao, Yiqing Liang, Yuanli Wang, Zilong Wang, Changzhi Zhou, David Heineman, Hange Liu, Harsh Trivedi, John Yang, Junhong Lin, Manish Shetty, Michael Yang, Nabil Omi, Negin Raoof, Shanda Li, Terry Yue Zhuo, Wuwei Lin, Yiwei Dai, Yuxin Wang, Wenhao Chai, Shang Zhou, Dariush Wahdany, Ziyu She, Jiaming Hu, Zhikang Dong, Yuxuan Zhu, Sasha Cui, Ahson Saiyed, Arinbjörn Kolbeinsson, Jesse Hu, Christopher Michael Rytting, Ryan Marten, Yixin Wang, Alex Dimakis, Andy Konwinski, and Ludwig Schmidt. Terminal-bench: Benchmarking agents on hard, realistic tasks in command line interfaces, 2026. 1 [42] Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave no context behind: Efficient infinite context transformers with infini-attention. arXiv preprint arXiv:2404.07143, 101, 2024. 5, 10 [43] OpenAI. Introducing swe-bench verified, 2024. 1 [44] OpenAI. Introducing gpt-5, 2025. 1 [45] Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc-Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th annual meeting of the association for computational linguistics (volume 1: Long papers), pages 1525–1534, 2016. 18 [46] Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin A Raffel, Leandro Von Werra, Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data at scale. Advances in Neural Information Processing Systems, 37:30811–30849, 2024. 2, 6 [47] Qwen Team. Qwen3-next: Towards ultimate training & inference efficiency, 2025. 1 [48] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018. 1 [49] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019. 1 [50] Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint arXiv:1911.05507, 2019. 18 [51] Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don’t know: Unanswerable questions for squad. arXiv preprint arXiv:1806.03822, 2018. 7, 18 [52] Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. Samba: Simple hybrid state space models for efficient unlimited context language modeling. arXiv preprint arXiv:2406.07522, 2024. 10 [53] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. In Proceedings of the AAAI Conference on Artificial Intelligence, pages 8732–8740, 2020. 18 [54] Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International conference on machine learning, pages 9355–9366. PMLR, 2021. 1, 2, 10, 19 [55] Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. arXiv preprint arXiv:2208.04933, 2022. 3 [56] Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder architectures for language models. Advances in Neural Information Processing Systems, 37:7339–7361, 2024. 6 [57] Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, et al. Learning to (learn at test time): Rnns with expressive hidden states. arXiv preprint arXiv:2407.04620, 2024. 1, 2, 3, 6, 10, 17, 21, 22, 23

13

[58] Arnuv Tandon, Karan Dalal, Xinhao Li, Daniel Koceja, Marcel Rød, Sam Buchanan, Xiaolong Wang, Jure Leskovec, Sanmi Koyejo, Tatsunori Hashimoto, Carlos Guestrin, Jed McCaleb, Yejin Choi, and Yu Sun. End-to-end test-time training for long context. arXiv:2512.23675, 2025. 1, 10, 22 [59] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023. 5, 7 [60] Harsh Trivedi, Niranjan Balasubramanian, Tushar Khot, and Ashish Sabharwal. Musique: Multihop questions via single-hop question composition. Transactions of the Association for Computational Linguistics, 10:539–554, 2022. 18 [61] Ke Alexander Wang, Jiaxin Shi, and Emily B Fox. Test-time regression: a unifying framework for designing sequence models with associative memory. arXiv preprint arXiv:2501.12352, 2025. 10 [62] Haoning Wu, Dongxu Li, Bei Chen, and Junnan Li. Longvideobench: A benchmark for long-context interleaved video-language understanding. Advances in Neural Information Processing Systems, 37: 28828–28857, 2024. 19 [63] Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, et al. Effective long-context scaling of foundation models. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pages 4643–4663, 2024. 8 [64] Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:2312.06635, 2023. 6 [65] Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint arXiv:2412.06464, 2024. 1, 4, 18 [66] Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. Advances in neural information processing systems, 37:115491–115522, 2024. 6, 7, 10, 18 [67] Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William Cohen, Ruslan Salakhutdinov, and Christopher D Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. In Proceedings of the 2018 conference on empirical methods in natural language processing, pages 2369–2380, 2018. 18 [68] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv:1905.07830, 2019. 18 [69] Tianyuan Zhang, Sai Bi, Yicong Hong, Kai Zhang, Fujun Luan, Songlin Yang, Kalyan Sunkavalli, William T Freeman, and Hao Tan. Test-time training done right. arXiv preprint arXiv:2505.23884, 2025. 1, 2, 3, 4, 5, 6, 7, 10, 17, 18, 19, 21, 22 [70] Yuanhan Zhang, Jinming Wu, Wei Li, Bo Li, Zejun Ma, Ziwei Liu, and Chunyuan Li. Video instruction tuning with synthetic data. arXiv preprint arXiv:2410.02713, 2024. 9, 18 [71] Ming Zhong, Da Yin, Tao Yu, Ahmad Zaidi, Mutethia Mutuma, Rahul Jha, Ahmed Hassan, Asli Celikyilmaz, Yang Liu, Xipeng Qiu, et al. Qmsum: A new benchmark for query-based multi-domain meeting summarization. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 5905–5921, 2021. 18

14

A

Detailed Derivation of the Parallel Kernel

In this section, we derive the closed-form expansion of the TTT update rule within a single chunk of length C. This formulation allows us to “jump” directly from the start of the chunk (t = 0) to the end (t = C) without materializing intermediate states. The chunk index r is omitted for simplicity. Recall that we consider general per-token updates with coupled L2 decay and momentum. Let the fast weights Wt ∈ Rd×d , momentum Mt ∈ Rd×d , and update direction (negative gradient) Gt ∈ Rd×d be matrix-valued. Let γt , βt ∈ (0, 1) be the scalar decay and momentum factors, and ηt > 0 be the learning rate. The recurrence is: Mt = βt Mt−1 + ηt Gt , A.1

Wt = γt Wt−1 + Mt .

(17)

Unrolling Momentum and Weights

Unrolling the momentum. By induction, the momentum Mt at step t is the sum of the decayed initial momentum at the chunk start and the history of gradients: Mt =

t Y

t  t   X Y βi M0 + ηj Gj βi .

i=1

j=1

(18)

i=j+1

Unrolling the Weights. Similarly, unrolling the weight update (Eq. 17) yields: Wt =

t Y

t  t   X Y γi W0 + Mj γi .

i=1

j=1

(19)

i=j+1

Insert (18) into (19) and separate the M0 and gradient parts: Wt =

t Y

j j  j t  Y t t X t   X   X Y Y Y βi M0 γq + γi W0 + ηk Gk βi γq . (20) j=1

i=1

|

{z

Decay term

}

i=1

|

q=j+1

{z

}

Momentum term

j=1 k=1

i=k+1

|

{z

q=j+1

}

Gradient term

Re-indexing the gradient contribution. Swap the order of summation in the last term of (20) by fixing the gradient index k and summing over j = k, . . . , t: j  t X X

j Y

ηk Gk

j=1 k=1

i=k+1

j t t  Y t  X  X Y γq = ηk Gk βi γq .

t Y

βi

q=j+1

k=1

j=k

(21)

q=j+1

i=k+1

This converts Eq. 20 to: Wt =

t Y

j j t t  Y t t  Y t     X X X Y Y γi W0 + βi M0 ηk Gk βi γq . (22) γq +

i=1

| A.2

j=1

{z

Decay term

}

|

i=1

q=j+1

k=1

{z

}

Momentum term

j=k

|

i=k+1

{z

Gradient term

q=j+1

}

Closed form at the final step per chunk (t = C)

To efficiently compute the final state WC , we define three auxiliary scalar sequences. Let the suffix products for decay γ and momentum β be: β̃t :=

C Y

γ̃t :=

βi ,

i=t+1

(with the empty product convention

C Y

γi ,

(23)

i=t+1

QC

C+1 (·) = 1). We further define the cumulative ratio sum Rt :

Rt :=

C X γ̃i i=t

β̃i

.

(24)

Since these terms depend only on the scalars β, γ, they can be pre-computed efficiently in parallel (e.g., via log-space cumulative sums). 15

Initial Weight Carry (carryw ). Directly from definition: carryw :=

C Y

γi = γ̃0

(25)

i=1

Initial Momentum Carry (carrym ). We rewrite the prefix product divided by the suffix product: j C Y Y
 βi = βi β̃j . i=1

Qj

i=1 βi as the total product

(26)

i=1

Substituting this into the M0 term in Eq. 22: j C  Y X j=1

C C  Y X
γ̃j βi M0 γ̃j = βi M0 , β̃ i=1 i=1 j=1 j | {z }

(27)

= R1

with carrym := β̃0 R1 . Gradient contributions. From (21) at t = C, each Gk is weighted by the scalar j C  Y X j=k

βi

i=k+1

(28)

C Y

C  X γ̃j γq = β̃k = β̃k Rk . β̃ q=j+1 j=k j

(29)

Closed form for WC . Combining these terms, the final weight matrix WC is computed as: WC = γ̃0 W0 + β̃0 R1 M0 +

C X

ηt β̃t Rt Gt .

(30)

t=1

Closed form for MC . The momentum recurrence is uncoupled from Wt , so its closed form follows directly from (18) by setting t = C: C  C C Y   X Y MC = βi M0 + ηt Gt βi . (31) t=1

i=1

i=t+1

Recognizing the suffix products defined in (23), β̃0 = collapses to the parallel kernel form MC = β̃0 M0 +

C X

QC

i=1 βi and β̃t

ηt β̃t Gt .

=

QC

i=t+1 βi , Eq. (31)

(32)

t=1

Unlike WC , the momentum kernel involves no cumulative ratio Rt : each gradient Gt is propagated to the chunk end through a single suffix product β̃t , since only one decay channel (β) acts on Mt . This formulation allows us to “jump” directly from t = 0 to t = C without materializing intermediate states. To ensure numerical stability, particularly with long chunk lengths C, we perform the cumulative product computations (β̃, γ̃) in log-space and convert back to linear space only for the final summation. A.3

Stability of the Cumulative Kernel Rt

We analyze the magnitude of the cumulative kernel Rt under two regimes that bound the practical operating range of our input-dependent dynamics: (i) near-unit retention (βt → 1 and γt → 1), as encouraged by the timescale parameterization (Eq. 13) for long-range memory, and (ii) aggressive decay (γt small). Near-unit retention regime. When βt → 1 and γt → 1 for all t, the suffix products satisfy β̃i → 1 and γ̃i → 1, so γ̃i /β̃i → 1 for every i. Substituting into the definition of Rt gives Rt →

C X

1 = C − t + 1 ≤ C.

i=t

16

(33)

The weight kernel KtW = ηt β̃t Rt is therefore bounded by KtW ≤ ηt · 1 · C. (34) Because ηt is sigmoid-gated and scaled by ηbase (Eq. 12), the per-token kernel contribution is bounded by ηbase · C. With our default ηbase = 10−2 and C = 512, this yields a worst-case bound of 5.12. The momentum kernel KtM = ηt β̃t is more tightly bounded by ηbase . Aggressive decay regime. When γt is small, the suffix product γ̃i decays geometrically in i, so the ratio γ̃i /β̃i decays accordingly. The sum Rt is then dominated by its first few terms and remains O(1). Coupled parameterization. The decay parameterization γt = 1 − ηt αt (Eq. 14) enforces a structural coupling: large updates (ηt large) automatically increase forgetting (γt smaller), preventing carry inflation across chunks. This analysis is consistent with the numerical verification reported in Sec. 3.2: over a 128-chunk trajectory (C = 512, 65,536 tokens), the relative ℓ2 deviation from the fully sequential reference is below 2 × 10−6 on both WC and MC . A.4

Single Backward Pass for Both Aggregates

Proposition 3.1 expresses both chunk-end states through a sum over per-token gradients Gt = −∇W Lt W : 0

C X

KtW Gt = −

t=1 C X

C X

KtW ∇W Lt W0 ,

(35)

KtM ∇W Lt W0 ,

(36)

t=1

KtM Gt = −

t=1

C X t=1

where the kernels KtW = ηt β̃t Rt and KtM = ηt β̃t differ only by the scalar factor Rt . Chain-rule decomposition. Each per-token loss has the form Lt = L(fW0 (kt ), vt ). Applying the chain rule, ⊤  ∂fW0 (kt ) ∂Lt gt , gt := . (37) ∇W Lt W = 0 ∂W ∂fW0 (kt ) ⊤ The Jacobian (∂fW0 (kt )/∂W) depends only on (W0 , kt ) and is independent of the scalar kernels. Shared backward, scalar re-weighting. Both aggregates therefore reduce to scalar-weighted C contractions over the same set of activation gradients {gt }C t=1 and per-token inputs {kt }t=1 , differing W M only in the scalar weight Kt versus Kt applied per token. A single autograd backward through the per-token losses produces {gt }; both aggregates are then constructed by re-using these intermediates with the appropriate scalar weighting. The marginal cost of producing MC in addition to WC is therefore one extra scalar-weighted contraction, not a second backward pass. This is consistent with the throughput measurements reported in Fig. 5, where E2 -TTT stays close to the throughput of LaCT [69] despite computing the exact closed-form kernel. A.5

Two Equivalent Views of the Chunk Dynamics

Within a chunk, all gradients Gt in our formulation are evaluated at the frozen W0 (Sec. 3.2). Under this constraint, the closed-form chunk-end states from Proposition 3.1 admit two complementary readings. Optimization-style view. Equation (7) reads as a token-wise state recurrence with momentum factor βt and L2 decay factor γt , both modulated per token. This is the framing we adopt throughout the paper, for consistency with the mini-batch and chunk-wise TTT literature [57, 8, 69]. It is genuinely descriptive at the inter-chunk level: MC is carried into the next chunk and accumulates momentum over the full sequence. Kernel-weighting view. Equivalently, the chunk-end states can be read as scalar-weighted aggregations of per-token gradients, C C X X W WC = γ̃0 W0 + β̃0 R1 M0 + Kt Gt , MC = β̃0 M0 + KtM Gt , t=1

t=1

17

in which (KtW , KtM ) act as learned, time-aware kernel weights over per-token gradient contributions to the chunk-end states. The two views are mathematically equivalent for the chunk-end pair (WC , MC ). They differ only in emphasis: the optimization-style view is the natural reading across chunk boundaries, where momentum is genuinely carried; the kernel-weighting view is the natural reading of the within-chunk aggregation, where the frozen-W0 constraint makes “momentum” and “decay” coincide with a learned per-token weighting of gradient contributions.

B

Experiment Continued

B.1

Training Details

Table 7 shows the training and model hyperparameters used to train the 340M and 1.3B parameter models. We train with an effective batch size of 256 with a sequence length of 2048 for 28,672 steps; this yields 15B tokens. The 1.3B models are trained with a slightly increased sequence length of 2240, which increases the training token count to 16B (for simplicity, in the main text, we refer to both as trained for 15B tokens). We implement our expressive chunk-wise TTT in plain PyTorch without custom Triton kernels. Training the 340M and 1.3B models takes about 132 and 348 total GPU hours on H100 GPUs, respectively. Following Zhang et al. [69], we apply L2 normalization to the fast-weight matrices at each chunk boundary as a standard stabilizer.

Table 7: Model hyperparameters. Model 340M 1.3B Number of layers Feedforward block multiplier Total hidden size Number of heads

24 4 1024 8

2048 16

Sequence length Effective Batch size Learning rate Warmup steps Minimum learning rate Max norm clipping Std. of weight initializers

2048 2240 256 1e−3 1024 0.1 1.0 0.02

TTT base learning rate ηbase

0.01

In our video understanding experiments, we inTTT momentum timescale τ 32 TTT decay strength αbase 0.1 stantiate E2 -TTT (parameterized by a SwiGLU MLP) as a parallel branch fused with the existing Qwen3VL self-attention layers. E2 -TTT employs its own query, key, and value projection layers. Following standard practice [1, 12], we gate the output of E2 -TTT with a learned vector ε ∈ Rd with tanh activation. All values in ε are initialized to 0.1, so the values in tanh(ε) are close to 0 at the beginning of fine-tuning. We train only the TTT parameters on a 43K-sample subset of LLaVA-Video-178K [70]. B.2

Evaluation Details

Commonsense Reasoning. Following prior work [21, 65, 28], we evaluate our model on multiple commonsense reasoning benchmarks: PIQA [10], HellaSwag [Hella.; 68], WinoGrande [Wino.; 53], ARC-easy (ARC-e) and ARC-challenge (ARC-c) [11], Wikitext [Wiki.; 40], and LAMBADA [LMB.; 45]. All evaluations are performed using lm-evaluation-harness [18]. In-context Retrieval. Following prior work [66, 3, 28], we focus on three tasks: FDA [3], SWDE [38], and SQuAD [51]. For SQuAD and SWDE, we use lm-evaluation-harness [18] for evaluation. For FDA, we follow Irie et al. [28] and use the evaluation script from Arora et al. [5]. Length Extrapolation. We first assess whether the models can maintain stable language modeling performance when processing sequences longer than their training context (2K tokens). Following Yang et al. [65], we conduct experiments on six diverse long-context benchmarks: PG19 [50], GovReport [27], QMSum [71], NarrativeQA [33], Qasper [14], and CodeParrot. To test effective capacity, we utilize the S-NIAH-1 (passkey retrieval) and S-NIAH-2 (numerical needle in haystack) tasks from RULER’s Needle-In-A-Haystack Single (NIAH-S) benchmark suite [24]. Finally, we evaluate on 14 real-world long-context tasks from Longbench [7], including: narrative comprehension (Narrative QA [33]), scientific understanding (QasperQA [14]), multi-hop reasoning (MultiField QA, HotpotQA [67], 2WikiMulti QA [23], Musique [60]), document summarization (GovReport [27], QMSum [71], MultiNews [17]), and various specialized tasks (TRec [35], Trivia QA [29], 18

SamSum [19], LCC [22], and RepoBench-P [37]). All evaluations, except for the stability analysis, are performed using lm-evaluation-harness [18]. Video Understanding. We evaluate on two standard video understanding benchmarks: VideoMMMU [25] and LongVideoBench [62]. Following the Qwen3VL [6] protocol, we employ a uniform sampling strategy to extract a fixed sequence of frames from each video. We configure the frame count and token budget to balance temporal range with spatial resolution: for VideoMMMU, we sample 512 frames with a maximum of 256 tokens per frame, and for LongVideoBench, 256 frames with 512 tokens. Hybrid Baselines. We benchmark against two competitive hybrid architectures: HQLT [28] and LaCT [69]. For HQLT, which combines DeltaNet [54] with window attention, we select the synchronous variant where both components process inputs simultaneously; this configuration is reported as the optimal setting by the authors. For LaCT, which fuses chunk-wise TTT (parameterized by a SwiGLU MLP) with window attention, we evaluate the variant utilizing Muon optimization for the TTT update rule, consistent with the best-performing configuration provided in the official implementation. B.3

Ablation Study

We validate the design of E2 -TTT along four axes. (i) The proposed closed-form kernel is the source of accuracy gains over chunk-wise baselines that average per-token factors, and finer chunk granularity narrows the residual gap to the ideal token-wise update (Table 8). (ii) Both branches of the hybrid architecture and the post-chunk normalization are individually necessary; removing any one of them collapses either retrieval or long-context performance (Table 9). (iii) The token-wise momentum and weight-decay terms recovered exactly by our kernel reduce perplexity over a no-momentum baseline (Table 10). (iv) The method is robust to the choice of base learning rate and base decay over a wide range (Table 11). Table 8: Effect of the closed-form kernel and chunk granularity. Trained on 15B tokens; 340M parameters. The averaged-factor baseline collapses per-token βt , γt into chunk-level scalars; our kernel preserves the full per-token dynamics in closed form. Configuration Averaged-factor baseline (C=512) Closed-form kernel, C=1024 Closed-form kernel, C=512

Wiki ppl ↓

LMB ppl ↓

Avg ↑

26.9 25.6 25.5

25.9 26.2 25.0

48.6 48.4 49.1

Effect of the closed-form kernel and chunk granularity (Table 8). We compare against a chunkwise baseline that averages the per-token decay and momentum factors into chunk-level scalars, following the simplification used in LaCT [69]. Replacing the averaged factors with our closed-form kernel reduces Wiki perplexity from 26.9 to 25.5 and improves average accuracy from 48.6% to 49.1% at C=512. We then sweep the chunk size: C=1024 already outperforms the averaged baseline (Wiki ppl 25.6, Avg 48.4%), and C=512 further improves both metrics. This trend matches the theoretical prediction that C → 1 recovers the token-wise update exactly; smaller chunks narrow the gap at the cost of hardware utilization, so we adopt C=512 as the efficiency–quality trade-off in the main experiments. Table 9: Component ablation at the 1.3B scale on S-NIAH-1 (accuracy ↑ at three context lengths) and LongBench (average score ↑). The full model is trained for 15B tokens at 2K context; ablations reuse identical training data and hyperparameters. Configuration SWA only (no TTT branch) TTT only (no SWA branch) Full, w/o post-chunk norm Full, w/o fusion gate Full E2 -TTT S WI GLU

S-NIAH-1 2K

S-NIAH-1 8K

S-NIAH-1 16K

LongBench

26.8 0.0 100.0 99.6 95.6

6.8 0.0 99.8 88.8 95.4

2.8 0.0 98.8 86.6 93.6

11.2 1.5 13.4 13.3 14.1

19

Branch contributions (Table 9). We isolate the role of each component of the hybrid architecture at the 1.3B scale, evaluating on Single-NIAH-1 at three context lengths and on LongBench. The results are unambiguous. Sliding-window attention alone reaches only 6.8% on S-NIAH-1 at 8K tokens, confirming that local attention cannot solve long-range retrieval. The TTT branch alone (no SWA) collapses to 0.0% across all S-NIAH lengths, because the chunk-wise output step cannot resolve within-chunk dependencies; this is exactly the “blind spot” our hybrid is designed to address. With both branches, S-NIAH-1 reaches 95.4% at 8K and 93.6% at 16K. Post-chunk normalization trades a small amount of synthetic-retrieval accuracy (99.8 versus 95.4 on S-NIAH-1 at 8K) for real-world long-context performance (LongBench 14.1 versus 13.4). Removing the fusion gate costs accuracy beyond 2K on both axes. Table 10: Effect of the inner-loop optimizer terms recovered by our kernel (340M, C=512, 15B tokens). The baseline retains only the input-dependent learning rate ηt . Inner-loop update rule

Wiki ppl ↓

LMB ppl ↓

Avg ↑

26.6 25.5

26.4 25.0

48.7 49.1

ηt only (no momentum, no decay) ηt + momentum + weight decay

Inner-loop momentum and weight decay (Table 10). We test whether the input-dependent momentum and decay terms contribute, by training two variants with C=512: a no-momentum/no-decay baseline (which retains only the input-dependent learning rate) and the full update rule. Adding momentum and weight decay reduces perplexity on both Wiki and LMB. Table 11: Sensitivity of E2 -TTT MLP (340M) to the base hyperparameters of the input-dependent dynamics. LMB ppl ↓

Avg ↑

Base learning rate ηbase (αbase =10 ) ηbase = 10−1 28.8 ηbase = 10−2 (default) 25.5 ηbase = 10−3 26.1

30.3 25.0 23.0

46.5 49.1 49.1

Base decay αbase (ηbase =10−2 ) αbase = 10−1 (default) 25.5 αbase = 10−2 25.8 αbase = 10−3 25.7

25.0 23.6 22.9

49.1 49.0 49.0

Configuration

Wiki ppl ↓ −1

Sensitivity to base hyperparameters (Table 11). The input-dependent dynamics expose a base learning rate ηbase and a base decay αbase that scale the per-token sigmoid gates. We sweep each over two orders of magnitude. Performance is essentially invariant to αbase in the tested range and stable for ηbase ∈ {10−3 , 10−2 }; only the aggressive ηbase =10−1 degrades training, consistent with classical optimization where overlarge step sizes destabilize updates. The defaults ηbase =10−2 and αbase =10−1 used throughout the paper are not finely tuned. B.4

Analysis of the Learned Coefficients

The ablations of Appendix B.3 remove a term from the update rule and retrain. A complementary question is whether the three linear heads of Eqs. (12)–(14) actually use their input-dependence once trained, or whether they instead settle near constants or silently compensate for one another. We answer this with two measurements that require no retraining: both run post-hoc on the single trained 1.3B E2 -TTT S WI GLU checkpoint used throughout the main experiments. The coefficients are not near-constant. We measure dispersion on the three positive quantities whose zero value means “no effect”: the step size ηt , the forgetting rate − log βt , and the decay rate 1 − γt = ηt αt . Over WikiText-2, the median coefficient of variation across all (layer, head) pairs is 0.95 for ηt , 0.46 for − log βt , and 1.29 for 1 − γt ; fewer than 1% of pairs fall below 0.1. Compounded over a chunk of C=512 tokens the spread Q is far from a rounding effect: the fast-weight Q carry t γt spans 0.73–1.00 and the momentum carry t βt spans a factor of 129 between its 5th and 95th percentile. This is the per-token variation that a chunk-averaged update rule discards. 20

Causal test: shuffling the coefficients along the token axis. A wide distribution alone does not establish that the values are aligned with the tokens that produced them. For each head we therefore permute its scalar along the token axis, independently per sequence and per head. The permutation preserves the marginal distribution exactly and destroys only the correspondence between a coefficient and its token, so any degradation is attributable to input-dependence alone. For α we permute the head output and recompute γt = 1 − ηt αt , which leaves ηt aligned. All variants are scored with the same lm-evaluation-harness invocation used for Table 2. Table 12: Shuffling each coefficient head along the token axis on the trained 1.3B E2 -TTT S WI GLU checkpoint. The permutation leaves the marginal distribution of each scalar unchanged and removes only its alignment with the input. Shuffled head

Wiki ppl ↓

∆ ppl

S-NIAH-1 16K ↑

— (baseline) η β α All three

19.94 20.39 20.05 19.99 20.54

— +0.45 +0.11 +0.05 +0.60

93.6 6.0 86.2 91.8 4.0

Every intervention hurts, and the ranking is informative. Shuffling η alone costs +0.45 perplexity and destroys retrieval outright (93.6 → 6.0 on S-NIAH-1 at 16K), confirming that the step size carries most of the write-selectivity of the inner loop. The joint intervention costs essentially the sum of the individual ones (+0.60 versus 0.45 + 0.11 + 0.05), which is the signature of three independent contributions: had the heads been substituting for one another, removing one alone would have been cheap, because the remaining two would absorb its role, and the joint cost would have exceeded the sum. Additivity is readable on perplexity only, since on S-NIAH the η intervention by itself already reaches the floor. The β and α interventions are gentler than the η one partly by construction: Eq. (14) couples the decay to the step size, so permuting only the α head leaves the η-driven component of γt aligned and intact. That the momentum and decay terms nonetheless matter in their own right is shown independently by Table 10: removing them from the update rule altogether costs 25.5 → 26.6 Wiki perplexity at 340M. Taken together, the two measurements support the premise behind the closed-form kernel of Sec. 3.2 — the per-token scalars are input-dependent, non-redundant, and worth aggregating exactly rather than collapsing to chunk averages. B.5

Throughput and Inference Cost Ours MLP

Ours SwiGLU

LaCT

Titans

23.7

20

91.8

20.0

91.2 79.5

80

70.5

17.6

15 10

9.2

8.9 7.6

TFLOPS

kTPS

63.0

60

58.0

40 20

5 1.0

0

340M

3.0

0

1.3B

340M

1.3B

Figure 5: Training throughput on a single H100 GPU at the 340M and 1.3B scales. To validate the hardware efficiency of our proposed method, we evaluate training throughput (measured in kilo-tokens per second, kTPS) and utilization (TFLOPS) on a single H100 64GB GPU. We test two model scales (340M and 1.3B parameters) on 2K sequences, using batch sizes of 4 and 2, respectively. We compare E2 -TTT against Titans [8] and LaCT [69]. All baselines are implemented in PyTorch under identical training conditions. As no official Titans implementation was available at the time of writing, we report a self-reproduction following the architecture described in Behrouz et al. [8]; we implemented both the scan-based formulation and the memory-optimized dual form [57] and 21

report the latter, which we found slightly faster in PyTorch. As shown in Figure 5, Titans is severely throughput-bound: at 340M it achieves less than 2 kTPS, and at 1.3B it goes out-of-memory. This is consistent with the structure of the dual form for nonlinear inner models such as TTT-MLP: each inner-model layer is reformulated as a masked attention-like computation (Appendix A.3, Sun et al. [57]), and this machinery lies on the outer-loop autograd path, inflating both activation memory and backward-pass cost. The dual form does eliminate per-token state materialization in principle, but for TTT-MLP the small inner mini-batch size (64 in our setting) combined with these masked-matmul overheads dominates throughput. These costs are absent under chunk-wise output (Eq. 6), where the per-chunk forward and backward operate on a single fast-weight state. In contrast, both E2 -TTT variants stay close to the throughput of the highly optimized LaCT baseline despite computing the exact closed-form kernel rather than a single averaged factor. The additional bookkeeping incurred by our kernel—log-space suffix products, ratio sums, and a scalar re-weighting that yields MC from the same backward pass that produces WC —is cheap relative to the fast-weight forward and backward themselves. Relative to LaCT, the exact kernel costs 15.6% throughput at 340M and only 3.3% at 1.3B for E2 -TTTMLP ; the overhead shrinks with scale because the d2 fast-weight forward and backward come to dominate the scalar bookkeeping. Exactness over per-token dynamics is therefore attainable at a small cost in training throughput. Table 13: Inference cost at 1.3B on a single H100 GPU, at batch sizes 1 and 8. Decode latency is the mean per-token time over 512 steps after warm-up (median of three runs). All recurrent models are flat in both latency and memory between 2K and 16K, while attention is not. Decode (ms/token) ↓

Peak mem. (GB) ↓

Model

2K

16K

2K

16K

Batch size 1 Transformer++ Titans LaCT E2 -TTTMLP E2 -TTTSwiGLU

14.8 56.1 36.8 33.8 34.0

15.0 56.2 36.5 33.6 33.9

3.25 3.41 3.33 3.43 3.62

6.10 3.51 3.42 3.52 3.71

Batch size 8 Transformer++ Titans LaCT E2 -TTTMLP E2 -TTTSwiGLU

14.8 56.5 38.2 33.7 34.3

59.9 56.6 38.8 33.8 34.4

7.50 6.40 7.10 6.86 8.08

29.90 6.40 7.10 6.86 8.08

Inference cost. Table 13 reports decode latency and peak memory at 1.3B. Attention is faster at short sequences thanks to its optimized kernels, but it is the only model whose cost grows with length (14.8 → 59.9 ms/token and 7.50 → 29.90 GB between 2K and 16K at batch 8), while all recurrent models are flat in both because their state is of constant size. Among TTT methods Titans is the slowest to decode, for the same reason it is slow to train. Both E2 -TTT variants decode slightly faster than LaCT at comparable memory: the per-chunk scalar bookkeeping is amortized over C tokens and is negligible next to the fast-weight forward. The one exception is the peak memory of E2 -TTTSwiGLU at batch 8 (8.08 against LaCT’s 7.10 GB), which comes from its three-matrix fast-weight network and the extra re-weighted aggregation (Algorithm 1). E2 -TTTMLP , whose fast-weight network is smaller, stays below LaCT at batch 8 (6.86 versus 7.10 GB) and within 0.1 GB of it at batch 1. B.6

Detailed S-NIAH Results

Section 5.3 of the main paper reports S-NIAH-1 and S-NIAH-2 results up to 16K for the four methods shown in Figure 4. Here we provide the corresponding numerical results, including two additional baselines not covered in the main figure: Mamba2 [13] and E2E-TTT [58]. E2E-TTT [58] is a recent chunk-wise TTT method that adopts a different inner-loop objective from the key–value binding formulation used in our work and in LaCT [69]. The original paper reports results only within the training context length and provides a JAX implementation. We re-implemented E2E-TTT in PyTorch following the architecture and update rule of the original paper, and trained it under our identical data, tokenizer, and optimization setup. As a sanity check, our re-implementation 22

Table 14: Zero-shot performance on the RULER S-NIAH benchmark suite for 1.3B models. Training context is 2K tokens; columns to the right of 2K are extrapolation. E2 -TTTSwiGLU retains a substantial margin over all baselines at every extrapolation length. Best result per column in bold. Method

16K

512

1K

S-NIAH-2 2K 4K

8K

16K

100.0 100.0 98.6 63.8 31.6 13.0 100.0 90.0 72.6 50.2 37.4 25.2 100.0 99.4 100.0 81.0 37.4 3.0 100.0 60.0 27.6 12.6 6.4 3.0 100.0 96.8 95.6 95.0 95.4 93.6 100.0 95.6 93.2 89.2 88.2 85.0

100.0 100.0 100.0 100.0 100.0 100.0

99.8 80.0 96.8 68.4 99.0 94.2

99.4 57.0 91.4 29.8 96.8 79.6

27.4 15.4 5.6 10.2 40.6 9.8

5.0 4.4 0.0 3.4 6.0 4.0

512 Mamba2 HQLT LaCT E2E-TTT E2 -TTTSwiGLU E2 -TTTMLP

1K

S-NIAH-1 2K 4K

8K

45.2 37.2 80.6 17.4 87.6 32.6

reaches ceiling accuracy with a sequence length of 512, consistent with other models. E2 -TTT retains a substantial margin over all baselines at every extrapolation length on both S-NIAH variants: at 16K on S-NIAH-1, E2 -TTTSwiGLU reaches 93.6% while all baselines fall below 26%; on S-NIAH-2, E2 -TTTSwiGLU remains the leading method through 16K, with all sub-quadratic baselines collapsing to below 30% from 8K onward. B.7

Training at a Larger Budget

To verify training feasibility at a larger budget, we additionally trained E2 -TTTSwiGLU at 1.3B parameters for 100B tokens (6.7× the budget of every other run in this paper) at a 4K training context, leaving the fast-weight network, the sliding-window size (512), the chunk size (512), the data source and the tokenizer unchanged. Table 15 reports the resulting checkpoint. Perplexity and commonsense accuracy improve substantially over the 15B-token model, and in-context retrieval rises from 43.6% to 58.1%. We stress that this run is intended only as a feasibility check. The two rows differ in both the token budget and the training context, so the gains cannot be attributed to the larger budget alone. What the run does show is that E2 -TTT trains stably at a budget 6.7× larger than the one used for the main experiments. Table 15: A 1.3B E2 -TTTSwiGLU checkpoint trained for 100B tokens at a 4K training context, shown next to the 15B-token checkpoint used throughout the main experiments. The Commonsense and Retrieval columns average the task sets of Tables 2 and 3; retrieval inputs are truncated to 2K and therefore fit inside the training context of both rows. Training tokens

Train ctx.

Perplexity ↓ Wiki. LMB.

15B 100B

2K 4K

19.9 16.0

15.8 10.0

Commonsense Avg. ↑

Retrieval Avg. ↑

53.6 60.1

43.6 58.1

Impact Statement and Limitations E2 -TTT advances the efficiency of long-context sequence modeling by parallelizing token-wise TestTime Training updates. As a foundational architecture for general-purpose language and multimodal models, it inherits the broad set of societal considerations associated with such models, but does not introduce risks specific to the proposed mechanism beyond those already present in existing TTT and linear-attention variants. We highlight two limitations of the current study. First, the chunk-wise output step (Eq. 6) applies a fixed fast-weight state to all tokens within a chunk, producing a within-chunk “blind spot” for finegrained causal dependencies; we mitigate this in practice through the sliding-window attention branch, and developing an efficient token-wise output step remains an interesting direction for future work. Second, following the standard evaluation protocol of recent TTT and linear-attention work [57, 28], our experiments are conducted at scales up to 1.3B parameters; extending E2 -TTT to larger model sizes is a natural direction for future work.

23

Related documents

• Rethinking Expressivity and Efficiency in Test-Time Training#677760
• Rethinking Expressivity and Efficiency in Test-Time Training#674550
• "It's about time: An analysis of student activities under remote learning"#694549
• Domestic Airlines - On Time Performance#398058
• The effects of odor and body posture on perceived duration#935042