Modular TTT: Rethinking Test-Time Training as Composable Modules

Abstract: Test-time training (TTT) views sequence modeling as an online learning problem in which fast weights are updated by an internal learning rule. Despite the growing number of TTT variants, existing approaches typically hard-code each variant separately, which makes it difficult to design new TTT methods and to isolate the role of each component. To address this, we propose Modular TTT, a framework that represents the inner learner as a directed acyclic graph and exposes the fast-weight network, loss function, learning rate, weight decay, and normalization as explicit design dimensions. Modular TTT automatically composes primitive-level train-view forward, train-view backward, and causal query-view rules into the full graph-level TTT computation, including the fast-weight state transition. Using Modular TTT, we systematically ablate the components of TTT and find that small learning-rate initialization, weight decay, and a single-layer nonlinearity improve performance, while MSE and inner-product losses perform similarly. Deeper fast-weight networks and normalization tend to hurt performance because they induce excessively large activations, while residual connections and gating provide little measurable benefit. Guided by these findings, we train the best resulting variant as 410M- and 1.45B-parameter models on 100B tokens, and observe training loss and benchmark performance comparable to Gated DeltaNet.

• The paper introduces a DAG-based framework that composes fast-weight maps, losses, normalization, gating, decay, and learning-rate rules through three causal computation passes.
• The paper finds that small learning-rate initialization near 10^-3, scalar decay, and a single SiLU or GELU activation improve performance, while deeper, residual, gated, or normalized learners often degrade it.
• The paper reports up to 2.2× training throughput for linear learners and 3.3× for MLP learners, with competitive language-modeling results but substantial weaknesses on precise long-context retrieval tasks.

Modular Test-Time Training as a Composable Design Framework

“Modular TTT: Rethinking Test-Time Training as Composable Modules” (2608.07110) addresses a methodological problem in test-time training (TTT): although TTT defines sequence modeling through online optimization of fast weights, existing variants are generally implemented as isolated, topology-specific systems. This implementation pattern makes it difficult to determine which architectural or optimization choices are responsible for observed improvements, and it imposes substantial derivation and engineering costs whenever a new inner learner is introduced.

The paper proposes Modular TTT, a framework in which the TTT inner learner is represented as a directed acyclic graph (DAG) of primitive operations. Fast-weight maps, nonlinearities, normalization, residual additions, multiplicative gates, learning-rate mechanisms, weight decay, and inner losses become independently configurable components. The central claim is that TTT should be treated not as a family of manually derived architectures, but as a compositional design space with reusable local update rules. The framework is evaluated through controlled ablations and large-scale language-modeling experiments at 410M and 1.45B parameters.

Motivation and Position within Efficient Sequence Modeling

TTT differs from conventional recurrent and state-space approaches by storing sequence information in the parameters of an inner model rather than in a fixed-dimensional hidden vector. At each position or chunk, the fast weights are updated using a self-supervised objective, typically by applying one or more gradient-based writes. This perspective connects TTT to fast-weight programming, linear attention, dynamic evaluation, and online optimization. The initial TTT formulation showed that expressive hidden states can be obtained by treating the recurrent state as the parameters of a learner (Sun et al., 2024).

Subsequent work expanded the design space through alternative fast-weight topologies, chunkwise implementations, and neural-memory objectives. Examples include “Test-Time Training Done Right” (Zhang et al., 29 May 2025), “Titans: Learning to Memorize at Test Time” (Behrouz et al., 2024), “ATLAS: Learning to Optimally Memorize the Context at Test Time” (Behrouz et al., 29 May 2025), and “Gated Delta Networks: Improving Mamba2 with Delta Rule” (Zhang et al., 29 May 2025). However, these systems often change several factors simultaneously. A new MLP-based learner may also introduce a different normalization scheme, loss, learning-rate parameterization, state-update rule, or kernel implementation. Consequently, comparisons across variants can conflate architectural capacity with optimization and systems effects.

Modular TTT is designed to separate these factors. The framework exposes five major dimensions: the fast-weight network, inner loss, learning-rate initialization and schedule, decay, and normalization. This factorization enables controlled experiments in which only one component changes while the backbone, training budget, chunk size, and implementation remain fixed.

The paper’s conceptual starting point is that linear learners, MLPs, and gated learners can all be treated as composable forms of TTT memory rather than as fundamentally separate systems.

Figure 1: Linear, MLP, and gated learners viewed as composable TTT memory forms.

The Three-Pass Computation

The key technical contribution is a general three-pass procedure for executing a graph-structured TTT learner. Let the inner learner be a DAG whose nodes represent primitive operations. The same graph is evaluated in three distinct computational views.

The first pass is the train-view forward. Keys are propagated through the learner graph using the current fast-weight state, producing a predicted target representation. If the learner maps keys $\mathbf K$ to $\hat{\mathbf V}$, the inner loss is computed against a target $\mathbf V$.

The second pass is the train-view backward. The gradient of the inner loss is propagated backward through the graph. Each primitive supplies a local backward rule that produces two kinds of information: gradients with respect to its inputs and a local fast-weight update. These signals are retained for the causal query computation.

The third pass is the query-view forward. Queries are propagated through the learner while incorporating the train-view activations, backward signals, and fast-weight updates. Unlike an ordinary forward pass, the query-view explicitly models the causal effect of all preceding writes. For a linear fast-weight node, the update can be expressed in chunkwise form using a lower-triangular interaction matrix, which enforces that token $t$ can only read writes generated at positions no later than $t$.

This decomposition is important because ordinary automatic differentiation is insufficient by itself. Automatic differentiation can generate local gradients for a specified train-view computation graph, but the causal query-view and fast-state transition require separately defined semantics. Modular TTT therefore registers train-view forward, train-view backward, and query-view operators for each primitive, then composes them according to the learner DAG.

Figure 2: Modular TTT executes graph-structured memory through shared train-view forward, train-view backward, and causal query-view passes.

The framework supports linear maps, elementwise activations, RMSNorm-like normalization, addition, multiplication, scalar or vector gating, and several inner losses. Because primitive rules are composed automatically in topological order, a new graph topology does not require a complete manual derivation of its global update equation. This is the principal distinction from earlier implementations in which each TTT variant maintained its own forward path, backward derivation, cache structure, and fused kernel.

The implementation also replaces nested autodifferentiation in the inner loop with analytic primitive-level backward operators. In microbenchmarks, the analytic linear backward operator is reported to be $1.65\times$ faster than an autodiff reference, while the analytic normalization backward operator is $2.62\times$ faster. For end-to-end training, Modular TTT achieves approximately $2.2\times$ the throughput of the official TTT implementation for a linear learner and approximately $3.3\times$ for an MLP learner.

Controlled Ablation of the TTT Design Space

The ablation study is conducted at 160M and 410M model scales using approximately 10B training tokens. The experimental design fixes the outer pre-normalized language-model backbone and varies the inner TTT components independently. The results establish a relatively narrow empirical frontier.

Loss functions

MSE and inner-product losses are essentially tied. At 160M, the reported validation losses are 3.0380 for MSE and 3.0383 for inner product; at 410M, they are 2.7949 and 2.7938, respectively. Five-seed experiments reinforce this conclusion: the mean losses for the 160M linear learner are 3.0381 for MSE and 3.0379 for inner product, while the corresponding 410M means are 2.7954 and 2.7934.

The paper attributes this similarity to the scale of the output gradient entering the fast-weight write. MSE preserves residual magnitude, while inner-product loss directly uses the target representation as the write signal. Both therefore retain information about the magnitude of the reconstruction or prediction signal. L1 removes residual magnitude by retaining only its sign, whereas RMSE normalizes the residual across the chunk. These losses perform substantially worse, supporting the claim that scale-bearing update signals are important for the evaluated TTT regime.

This result is technically significant because it weakens the assumption that sophisticated reconstruction losses are necessary for TTT. The loss choice matters, but the competitive region appears to be defined primarily by whether the induced gradient preserves useful update scale rather than by a unique objective formulation.

Learning-rate initialization

The most consistent optimization result is the importance of small learning-rate initialization. The standard parameterization initializes the learned inner learning rate near $1$, whereas the preferred configuration initializes it near $\hat{\mathbf V}$0. At 410M with MSE and scalar decay, this change reduces validation loss from 3.0820 to 2.7949. Without decay, the corresponding reduction is from 3.4036 to 2.9343.

The explanation is both dynamical and spectral. For an MSE update, the homogeneous component of the fast-weight transition contains a term of the form

$\hat{\mathbf V}$1

If the learning rate is too large relative to the key covariance, this operator can acquire eigenvalues outside the unit disk, causing amplification across the inner update. With a chunk size of 256 and an initial learning rate of $\hat{\mathbf V}$2, the nominal product of chunk length and learning rate is only $\hat{\mathbf V}$3, which controls the initial write magnitude without permanently fixing the learned learning rate. The outer model can subsequently adapt the learning-rate predictor.

This finding is one of the paper’s strongest practical recommendations: the inner learning rate should not be initialized according to the scale conventions of ordinary linear-attention accumulation. TTT writes accumulate gradients rather than values, so an initialization near unity can produce an excessively large chunk update.

Decay and forgetting

Decay provides a clear quality-efficiency trade-off. No decay is consistently inferior, scalar decay recovers most of the gain at negligible computational cost, and vector decay produces the lowest validation loss but significantly reduces efficiency.

At 410M with MSE, scalar decay obtains a validation loss of 2.7949 at approximately 34,978 training tokens per accelerator per second, while vector decay reaches 2.7821 at approximately 27,074 tokens per accelerator per second. Vector decay also increases peak memory from roughly 37.3 GB for scalar decay to 40.8 GB. The paper characterizes this as approximately a 25% throughput reduction and about 3 GB of additional memory.

The interpretation is that decay supplies a controlled forgetting mechanism. Without decay, fast-weight contributions accumulate indefinitely and stale context cannot be attenuated. Scalar decay applies a global contraction and captures most of the benefit of feature-selective vector decay. The authors therefore use scalar decay in subsequent experiments.

Pointwise nonlinearities and normalization

A single non-fast-weight nonlinearity after the fast linear map improves performance. At 410M with MSE, Linear + GELU obtains a loss of 2.7844 and Linear + SiLU obtains 2.7849, compared with 2.7949 for the linear baseline. The corresponding throughput reduction is modest but measurable: Linear achieves approximately 34,978 tokens per accelerator per second, whereas Linear + SiLU achieves approximately 32,129.

The paper interprets this improvement as coordinatewise modulation of the write. For a linear fast map followed by an activation, the gradient entering the fast weight is multiplied by the activation derivative. The activation therefore gates the memory write while remaining outside the fast-weight state itself. Bounded derivatives also avoid some of the activation amplification associated with normalization.

Normalization behaves less robustly. Linear + Norm reaches a loss of 2.8192 at 410M under MSE, worse than both activation-based variants. The authors attribute this behavior partly to the inverse standard-deviation term in RMSNorm’s backward pass, which can amplify gradients when the input norm is small. Additional experiments show strong sensitivity to the normalization $\hat{\mathbf V}$4 parameter. Increasing $\hat{\mathbf V}$5 can stabilize certain configurations, but normalization does not provide a consistent advantage over simple pointwise nonlinearities.

Deep, residual, and gated fast learners

The paper makes a deliberately negative claim: in the evaluated one-step TTT regime, increasing the depth or structural complexity of the fast learner does not improve over a shallow learner and often degrades performance.

At 160M, the shallow Linear-SiLU reference achieves validation loss 3.0205. The best completed product-form Linear-Linear variant reaches 3.1265, an absolute degradation of 0.1060. The best activation-only deeper variant, Linear-SiLU-Linear-SiLU, obtains 3.1156. A stabilized norm-containing learner reaches 3.1144. Residual and gated variants are also inferior: a residual learner obtains 3.1643, while a stabilized SwiGLU + Norm learner obtains 3.2041.

The paper provides a structural explanation rather than attributing the result solely to failed hyperparameter tuning. For a two-factor fast learner,

$\hat{\mathbf V}$6

the update of each factor depends on the activations and downstream weights of the other factor. The effective update of the product therefore contains additional preconditioning and second-order-like coupling terms. Two factorizations that represent the same effective matrix can induce different update dynamics because rescaling one factor and inversely rescaling the other changes the factor-level gradients.

This creates a gauge-like optimization degree of freedom: the system must learn both a useful effective memory function and a factorization whose induced one-step update is stable. A shallow linear learner does not have this additional burden. The result is consistent with classical analyses of deep matrix factorization, but the paper’s contribution is to connect that factorization geometry directly to online fast-weight updates.

The zero-initialization analysis provides a limiting example. A single zero-initialized linear fast weight can receive a nonzero gradient, whereas a product of two zero-initialized fast weights can remain trapped because every factor’s gradient contains a zero upstream or downstream factor. Gaussian initialization avoids this particular failure mode, but it does not eliminate the broader factor-coupling problem.

Long-Context Behavior and Scaling

Following the ablations, the authors train selected variants at 410M and 1.45B scales for approximately 100B tokens using 4K training contexts and chunk size 256. The main comparisons include MSE and inner-product losses, linear and Linear-SiLU learners, scalar decay, and small-learning-rate initialization.

At 410M, the best Modular TTT multiple-choice average is 50.99 for the MSE Linear-SiLU model, compared with 50.45 for GDN and 50.63 for LLaMA in the reported comparison. Its average perplexity across WikiText-103 and LAMBADA is 22.43, while GDN obtains 21.81. At 1.45B, the strongest Modular TTT multiple-choice average is 56.44 for the inner-product linear model, close to GDN’s 56.30 and above LLaMA’s 55.67 in the reported evaluation. Its average perplexity is 14.05, compared with 13.70 for GDN and 14.02 for LLaMA.

These results support the paper’s claim that a shallow Modular TTT model can achieve competitive language-modeling and multiple-choice performance against strong recurrent baselines. They do not, however, establish parity across all capabilities. Containment-style tasks remain substantially weaker. At 1.45B, the best Modular TTT containment average is 35.46, compared with 36.97 for GDN and 65.20 for LLaMA.

The long-context retrieval results are more restrictive. On the RULER single-needle family at 1.45B and 8K context, the best Modular TTT score is 21.07, compared with 18.07 for GDN and 92.13 for LLaMA. At 410M and 8K, the strongest Modular TTT single-needle score is 14.07, while LLaMA reaches 51.87. Multi-key and multi-query tasks are even more difficult, with several Modular TTT variants approaching very low accuracy at longer contexts.

Figure 3: Per-token loss remains stable around and beyond the 4K training context for representative 410M shallow Modular TTT variants.

The per-token loss analysis nevertheless indicates that the selected shallow variants do not exhibit an abrupt degradation immediately after the 4K training context. This suggests that the models maintain stable sequential processing beyond the training window, but stability should not be confused with reliable content-addressable retrieval. A model can maintain bounded loss while still failing to preserve the precise information required by needle-in-a-haystack tasks.

Systems Implications

Modular TTT has two distinct systems contributions. First, it reduces the cost of implementing new inner learners by replacing topology-specific derivations with a library of primitive operators. Second, it improves execution efficiency through analytic backward rules and fused chunkwise computation.

The reported throughput gains over the official TTT implementation are substantial. For TTT-Linear at 160M, throughput increases from 43,397.5 to 93,366.9 tokens per second, approximately $\hat{\mathbf V}$7. For TTT-MLP, it increases from 21,336.9 to 71,101.1 tokens per second, approximately $\hat{\mathbf V}$8. The improvement is particularly large for the MLP topology because nested autodifferentiation and topology-specific backward computation impose greater overhead.

The inference profile is more nuanced. GDN has the highest prefill and decoding throughput in the reported comparisons. Modular TTT is faster than the evaluated LaCT implementation and uses relatively low peak memory, but it remains slower than GDN. At 410M, for example, GDN reaches approximately 400.7 decode tokens per second, whereas the Modular TTT linear variants reach approximately 219 tokens per second. The discrepancy arises from the chunkwise query update: within an incomplete chunk, the implementation must account for all causal writes accumulated up to the current position, whereas a recurrent delta-rule model can update its state token by token with a simpler recurrence.

This distinction matters for deployment. Modular TTT is particularly attractive for training and high-throughput prefill, where chunkwise parallelism can be exploited. Autoregressive decoding requires further kernel and state-update optimization before the framework can compete with the most efficient recurrent alternatives.

Theoretical and Practical Implications

The paper’s main theoretical implication is that TTT architectures should be analyzed at the level of update geometry, not only representational capacity. A deeper fast learner is more expressive as a static function, but its one-step online optimization dynamics may be substantially less favorable. Factorization introduces conditioning, scale non-identifiability, and cross-factor gradient coupling. Consequently, increasing inner-model depth can reduce effective memory quality even when the outer architecture has sufficient capacity.

The empirical results also challenge several common design intuitions. Normalization, residual connections, and gating—mechanisms that are often beneficial in conventional deep networks—do not automatically improve fast-weight memories. Their effect depends on how they transform the train-view gradients and causal state transitions. In this setting, a single pointwise activation after a shallow fast map is more reliable than deeper nonlinear or gated inner networks.

Practically, the framework suggests a disciplined recipe for TTT LLMs:

1. initialize learned inner learning rates near $\hat{\mathbf V}$9 rather than near unity;
2. use scalar decay as a low-cost forgetting mechanism;
3. prefer MSE or inner-product losses;
4. use a shallow linear fast learner with an optional SiLU or GELU;
5. treat normalization and fast-learner depth as instability-sensitive choices rather than default improvements;
6. use analytic primitive-level backward rules and chunkwise fused kernels for training efficiency.

The limitations are equally important. The study evaluates autoregressive language modeling under a particular one-step update rule, chunking strategy, tokenizer, and backbone family. The negative results for deeper and gated learners do not imply that such architectures are intrinsically unsuitable for TTT. They may become useful with multiple inner steps, momentum, learned preconditioning, optimizer-state memory, alternative chunk schedules, Muon-like update normalization, or hybrid attention pathways. The paper explicitly notes that several system-specific refinements used by other TTT systems are outside the current fused graph abstraction.

Future Directions

The most immediate direction is to extend the primitive interface beyond first-order gradient writes. Momentum, adaptive preconditioning, second-order approximations, and multi-step inner optimization could provide deeper learners with better-conditioned updates. A compositional framework for optimizer state would allow the fast learner to represent not only parameters but also auxiliary statistics such as moments, confidence estimates, or retrieval-specific keys.

A second direction concerns retrieval. The large gap on RULER indicates that stable loss and competitive multiple-choice accuracy do not guarantee precise long-context recall. Future TTT memories may require explicit content-addressable mechanisms, sparse writes, write arbitration, or hybrid fast states that separate semantic compression from exact retrieval. Vector decay and learned write gates are natural starting points, but the results of this paper suggest that these mechanisms must be evaluated jointly with update stability and system cost.

A third direction is automated architecture search over learner DAGs. Modular TTT provides the representation needed for such search, but the current study performs manually designed ablations. Differentiable or Bayesian search over graph topology, primitive placement, loss functions, decay, and learning-rate schedules could reveal combinations that are difficult to identify through isolated sweeps. Such search should include computational constraints because the best validation-loss configuration—vector decay, for example—is not the best throughput-quality compromise.

Finally, deployment-oriented work should focus on reducing decode overhead. A truly competitive TTT system will need a state representation and kernel schedule that preserve chunkwise training efficiency while offering token-level decoding comparable to recurrent delta-rule models. This may require separate prefill and decode algorithms, partially materialized chunk states, or hardware-specific recurrence kernels.

Conclusion

“Modular TTT: Rethinking Test-Time Training as Composable Modules” (2608.07110) introduces a systematic abstraction for designing and analyzing test-time-trained sequence models. Its central mechanism is a DAG-based inner learner whose primitive train-view forward, train-view backward, and query-view rules are automatically composed into a causal fast-weight state transition.

The experiments identify a compact and practically useful design frontier: small learning-rate initialization, scalar decay, and a single pointwise nonlinearity improve performance; MSE and inner-product losses are comparable; normalization is unstable; and deeper, residual, or gated fast learners do not improve over shallow alternatives in the evaluated one-step setting. At 410M and 1.45B parameters, the resulting models achieve language-modeling and multiple-choice performance comparable to GDN while retaining substantial training-throughput advantages over the official TTT implementation.

The broader significance lies less in proposing another isolated TTT topology than in establishing a framework in which TTT design choices can be separated, recombined, benchmarked, and analyzed through their induced update dynamics. Further progress will depend on extending this modularity to richer optimization states, retrieval-oriented memory mechanisms, and inference-efficient state transitions.