PaperlayerView exact revision on arXiv

Source-linked AI summary

Modular TTT: Rethinking Test-Time Training as Composable Modules

Bohao Tang, Zhen Qin, Yuqi Pan, Zheng Li, Pengfei Liu, Ya Zhang

arXiv:2608.07110v1cs.LGcs.CL

TL;DR

Existing TTT variants hard-code their components, leaving the design space and individual module contributions difficult to study. The paper introduces Modular TTT, a DAG-based framework that composes primitive update rules and enables systematic ablations. The resulting choices yield scaled models with training loss and benchmark performance comparable to Gated DeltaNet.

  • Problem

    Existing TTT variants hard-code multiple components together, making new methods difficult to develop systematically and individual component roles difficult to isolate.

  • Method

    Modular TTT represents the inner learner as a directed acyclic graph and automatically composes primitive train-view forward, train-view backward, and causal query-view rules, including fast-weight state transitions.

  • Results

    Systematic ablations favor small learning-rate initialization, weight decay, and a single-layer nonlinearity; scaled variants achieve training loss and benchmark performance comparable to Gated DeltaNet.

  • Takeaways & Limitations

    Modular TTT provides a unified space for expressing, implementing, and systematically analyzing TTT variants through independently controllable design dimensions.

  • Takeaways & Limitations

    The study is limited to autoregressive language modeling under the reported training budgets and model scales, and precise long-context recall remains limited for shallow variants.

Abstract

from arXiv · show

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.

1 Introduction

Modular TTT addresses the poorly understood design space of test-time training by representing inner learners as composable graphs with independently controllable components. Systematic ablations identify effective choices and produce scaled models comparable to Gated DeltaNet.

  • Motivation: Test-time training casts sequence modeling as online learning, updating fast weights with an internal learning rule as sequences are processed.This expands state-update design beyond fixed recurrences and explicit caches.
  • Motivation: Existing TTT variants hard-code multiple components together, obscuring which design choices cause observed behavior and hindering systematic method development.The coupled choices include the fast-weight network, loss, learning rate, weight decay, and normalization.
  • Framework: Modular TTT represents the inner learner as a directed acyclic graph and independently exposes the fast-weight network, loss, learning rate, weight decay, and normalization.It composes primitive train-view forward, train-view backward, and causal query-view rules into the graph-level computation and state transition.
  • Findings: 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, whereas residual connections and gating provide little measurable benefit.
  • Scale-up: At 410M and 1.45B parameters, the selected Modular TTT variant achieves training loss and benchmark performance comparable to Gated DeltaNet.The models were trained on 100B tokens.

2 Related Work

Efficient sequence-modeling approaches preserve long-context ability while reducing sequence-length costs, but they differ in how they represent and update state. TTT uses online optimization of fast weights and can instantiate several other model families.

  • Overview: Efficient sequence-modeling work broadly includes recurrent models, linear attention, state space models, and test-time training.These families aim to reduce computation and memory costs that grow with sequence length while preserving long-context modeling.
  • Recurrent models: Recurrent models compress historical context into a fixed-size hidden state updated through explicit recurrence rules.Gating and nonlinear transformations can increase expressiveness, while removing nonlinearities can improve efficiency with competitive performance.
  • Linear attention: Linear attention aggregates key-value pairs using associative matrix multiplication, avoiding explicit construction of a quadratic attention matrix.Chunkwise recurrent computation improves GPU utilization while retaining linear sequence-length complexity in causal settings.
  • State space models: State space models represent long-range dependencies through structured continuous- or discrete-time state transitions.Specialized initialization, structured parameterization, content-dependent decay, and hardware-friendly scans are emphasized in representative methods.
  • Test-time training: TTT replaces a fixed-form hidden state with fast weights of a learnable model, updated through optimization steps on a self-supervised objective.With suitable losses and components, this perspective can instantiate Linear Attention, Linear RNN, and SSM variants.

3 Method

Modular TTT expresses a learner as a computation graph whose local train-view and query-view rules are automatically composed. The three-pass procedure produces updates from a training view and applies them causally to query inputs.

  • Train-view forward: The train-view forward executes graph nodes in topological order to compute the learner output from input keys.For an m-layer module, nonlinear activation functions are associated with the indexed inner-learner layers.
  • Train-view backward: The train-view backward computes the inner loss, intermediate gradients, local parameter updates, and backward signals in reverse topological order.The resulting activations and gradients are retained for the subsequent query-view computation.
  • Query-view forward: The query-view forward processes query inputs causally using the train-view activations, gradients, and parameter updates to produce the TTT module output.Its causal readout uses the lower-triangular structure of the query computation.
  • Composition: Primitive-level train-view forward, backward, and query-view rules are composed automatically across the graph, avoiding topology-specific global update derivations.Automatic differentiation supplies local backward signals, while registered primitive rules define the causal readout and fast-weight state transition.
  • Graph representation: A TTT learner can be represented as a directed acyclic graph whose nodes are primitive operations and whose edges encode tensor dependencies.Primitives include linear maps, elementwise nonlinearities, residual additions, and normalization layers.

4 Experiments

Experiments use Modular TTT to systematically ablate design choices, benchmark implementation efficiency, and evaluate shortlisted variants at larger scale. The results favor small learning rates, decay, shallow nonlinear learners, and selected efficiency–performance trade-offs, while deeper learners and normalization are less reliable.

  • Experimental setup: 160M- and 410M-scale ablations vary loss, learning-rate initialization, decay, and nonlinearity under matched experimental settings.The first stage trains models for 10B tokens; larger-scale studies train shortlisted variants for 100B tokens and compare them with representative baselines.
  • Loss ablation: MSE and inner-product losses are competitive, whereas L1 and RMSE perform substantially worse.MSE preserves residual magnitude, while inner product uses the target value as the write signal; both maintain informative update scales.
  • Learning-rate initialization: η0 = 10^-3 small-learning-rate initialization consistently improves stability and final performance over initialization near 1.Excessively large learning rates can make the update spectrum unstable when eigenvalues exceed magnitude 1.
  • Decay ablation: Scalar decay provides most of vector decay’s quality gain at negligible efficiency cost, while vector decay has the lowest loss but roughly 25% lower throughput and about 3 GB more peak memory.No decay is weaker because past fast-weight contributions accumulate without attenuation; scalar decay is therefore used in subsequent experiments.
  • Nonlinearity and normalization: GELU and SiLU improve performance, while normalization is less stable and can occasionally degrade performance.Pointwise activations provide the best quality–efficiency trade-off; bounded derivatives also avoid excessively large activations, unlike some normalization behavior.
  • Deep memory: Deeper fast-weight learners do not surpass shallow variants, and residual or gated variants provide little measurable benefit in the evaluated settings.Deep learners can diverge or remain behind the shallow frontier, consistent with the difficulty of optimizing both effective weights and their factorization.
  • Efficiency: 1.65× and 2.62× speedups are achieved by analytic Linear and Norm operators, while Modular TTT improves end-to-end throughput by 2.2×–3.3× over the official implementation.The analytic Norm operator also reduces peak memory from 31.3 to 19.3 MB under the evaluated software configuration.
  • Scale-up evaluation: At 410M and 1.45B parameters, shortlisted variants remain competitive with recurrent baselines, but containment-style tasks and precise long-context recall remain challenging.At 1.45B, the inner-product linear variant is close to GDN on average multiple-choice accuracy; LLaMA remains substantially stronger on RULER, especially at 8k.

5 Conclusion

Modular TTT factorizes the inner learner into independently controllable design dimensions and composes local rules over a fast-weight DAG. This reduces the cost of constructing and comparing variants while enabling systematic analysis and competitive performance at scale.

  • Conclusion: Modular TTT represents the fast-weight network as a DAG of registered primitives with independently controllable design dimensions.The framework exposes local train-view forward, train-view backward, and causal query-view rules.
  • Conclusion: The shared composition procedure avoids topology-specific global update derivations when recombining existing primitives.This makes constructing new TTT variants less costly.
  • Conclusion: Modular TTT enables systematic analysis of which TTT components provide consistent gains and which provide little benefit.The framework turns TTT design into a modular space for controlled exploration.
  • Conclusion: The resulting best variant achieves performance comparable to GDN at 410M and 1.45B scales.The comparison is reported at both model scales.

A Method details

Modular TTT executes a learner DAG using registered primitive rules for training, backward propagation, querying, and fast-state updates. Its implementation supports configurable losses, normalization, learning-rate and decay mechanisms, and causal chunkwise updates.

  • Graph execution: The implementation represents graph memory with ordered nodes and a designated output, where linear nodes carry fast weights.Activation, addition, multiplication, and normalization nodes transform activations or gradients.
  • Graph execution: The same graph is executed in train-view forward, train-view backward, and query-view forward passes.Registered primitive rules define the local behavior in each view.
  • Losses: MSE and inner-product losses preserve chunk-level scale information, whereas L1 removes magnitude and RMSE normalizes it.These rules are used in the loss-choice analysis.
  • Query and state update: The causal query-view readout uses an inclusive lower-triangular operator to enforce within-chunk causality.The end-of-chunk state is updated from the chunkwise fast-weight computation.
  • Learning rate and decay: Learning-rate factors are injected at fast nodes through scaled keys, while mean scaling divides the output gradient by chunk length.With s = 1, update scale is controlled by the learned learning-rate factor and its initialization.
  • Initialization: Zero initialization degrades performance across the tested configurations relative to the official Gaussian fast-weight initialization.The comparison covers the flash and official initialization choices.

B.1 Model architecture

The model uses a pre-normalization residual backbone with RMSNorm, a graph-memory token mixer, and a bias-free SiLU GLU channel mixer. Architecture and initialization vary with model scale and use distinct parameterization choices for Modular TTT projections.

  • Model architecture: Each block applies RMSNorm, a token mixer, a residual connection, RMSNorm, and a bias-free GLU channel mixer with SiLU activation.The channel mixer uses gate, up, and down projections.
  • Token mixer: The Modular TTT token mixer uses graph memory with query, key, value, and output projections.Large-scale variants use L2-normalized keys and leave queries unnormalized.
  • Configuration: The architecture hyperparameters include hidden dimension, layer count, head count, per-head dimension, and GLU inner width.These quantities are identified in the architecture table.
  • Initialization: Standard linear layers use truncated-normal initialization with standard deviation 0.02, while Modular TTT query, key, and value projections use Xavier uniform initialization with gain 0.01.The output projection is zero-initialized by residual rescaling, and fast weights use the official TTT initialization.

B.2 Training details

Training uses next-token prediction on a large-scale English corpus with fixed-length packed sequences and scale-dependent budgets. The analysis studies how modular choices affect fast-weight update scale, direction, and conditioning.

  • Training setup: All language-model runs use GPT-2 BPE tokenization, a 50,257-token vocabulary, and the next-token prediction objective.The standard training split is used.
  • Training setup: Ablation runs use sequence length 2048, while scale-up runs use sequence length 4096.Packed sequences may cross document boundaries because no document-boundary attention mask is inserted.
  • Training budget: The two ablation settings use 8 × 32 × 1 × 2048 × 20,000 = 10.5B tokens, while scale-up settings use 16 × 8 × 2 × 4096 × 100,000 = 104.9B tokens.The budgets follow the devices, batch, accumulation, context, and steps configuration.
  • Evaluation: Evaluation covers perplexity, multiple-choice accuracy, and containment-style tasks using lm-eval-harness.Group averages are arithmetic means over the listed metrics.
  • Update analysis: The update analysis tracks how each modular choice changes fast-weight write scale, direction, or conditioning.The analysis proceeds through learning-rate initialization, losses, decay, nonlinear modules, normalization, and multilayer learners.
  • Update scale: For one linear fast primitive, the update norm scales as ∥∆W∥F = O(ηc/s).Mean scaling and small-learning-rate initialization control the initial chunk-update scale in different ways.

C.2 Small-lr init (η0 = 10−3)

Modular TTT predicts a per-token learning rate from the outer representation, with small-lr initialization controlling update scale at the start of training. This reduces instability risk while allowing the learned rate to adapt during training.

  • The learning rate η is predicted from the outer representation rather than fixed as a scalar.
  • Standard initialization without mean scaling produces O(c) initial chunk scale and is empirically worse than small-lr initialization.
  • Small-lr initialization lowers the initial spectral scale of the MSE update, reducing the risk of eigenvalues with magnitude greater than one.
  • Unlike mean scaling, small-lr initialization only sets the starting value, allowing the predicted learning rate to grow or shrink per token during outer-loop training.
  • η0 = 10−3 gives p = 5 × 10−4 and b ≈ −7.60 under the specified sigmoid parameterization.
  • With ηinit = 10−3 and c = 256, the scale factor is c · 10−3 = 0.256.

C.3 Loss analysis

The loss and fast-learner analyses distinguish scale-preserving choices from weaker alternatives and show that deeper product-form learners alter the inner-update geometry. The controlled findings favor MSE or inner-product loss, shallow nonlinear learners, and decay over deeper, normalized, residual, or gated designs.

  • Loss analysis: MSE and inner-product loss are competitive, whereas L1 and RMSE are weaker because their gradients preserve less residual scale information.
  • Loss analysis: L1 retains the residual sign pattern while RMSE normalizes its gradient direction, so both remove magnitude information from the memory write.
  • Fast-learner geometry: Gaussian initialization avoids the zero-gradient boundary case, but deeper learners remain factor-coupled under one-step Modular TTT updates.
  • Fast-learner geometry: A shallow linear learner has a convex-quadratic MSE update, while deeper learners couple factor updates to upstream activations and downstream factors.
  • Fast-learner geometry: All-zero initialization gives zero gradients for product-form learners with depth L ≥2, while a single linear learner can still receive a nonzero update.
  • Component findings: Scalar decay improves over no decay, while simple GELU and SiLU activations improve performance and normalization remains mixed or unstable.

D.4 Multilayer and gated memories

In the controlled one-step Modular TTT setting, deeper, residual, and gated fast memories do not improve over the shallow Linear-SiLU reference. Initialization and normalization affect stability, but the strongest results remain shallow activation-only learners.

  • Multilayer memories: The best deeper row remains worse than the shallow Linear-SiLU reference by 0.0939 validation loss.
  • Multilayer memories: All matched two-layer activation and normalization variants remain behind the shallow reference, although activation placement provides small gains.
  • Robustness: Across chunk sizes and model scales, Linear-GELU and Linear-SiLU are strongest, while deeper Linear-SiLU-Linear-Norm remains behind every single-layer learner.
  • Residual and gated memories: Residual and gated extensions do not improve the controlled setting; residual rows trail matched non-residual counterparts, and SwiGLU requires additional stabilization.
  • Initialization: Zero initialization works for single linear and shallow Linear-SiLU memories but fails for deeper and gated product-form memories; Gaussian initialization avoids this failure mode.
  • Inference efficiency: Modular TTT updates fast-weight state only at chunk boundaries during decoding, adding overhead within incomplete chunks relative to GDN’s token-wise recurrent update.
  • Inference efficiency: GDN has the highest prefill and decode throughput, while every evaluated Modular TTT configuration uses less peak allocated memory than GDN and LaCT.

D.7 RULER evaluation details

The RULER evaluation covers several needle-in-a-haystack task families across 1k–8k contexts. Modular TTT variants are competitive with recurrent baselines in some settings but lag LLaMA on precise retrieval, especially at longer and harder contexts.

  • The complete evaluation includes niah_multikey, niah_single, niah_multiquery, and niah_multivalue task families at 1k–8k contexts.
  • The niah_single family provides a compact view of direct single-key retrieval, while niah_multikey family rows average three corresponding templates.
  • Modular TTT variants are competitive with recurrent baselines on some short- and mid-context settings but do not match LLaMA on precise retrieval.
  • The retrieval gap is largest at 4k–8k contexts and on harder multi-key templates, where several models are near zero.
  • At 410M, inner-product loss is often stronger for explicit retrieval, while SiLU variants are less consistently beneficial than in language modeling.

E Limitations

The study evaluates Modular TTT only within autoregressive language modeling at the reported scales and training budgets, while leaving broader inner learners, update schedules, and retrieval settings open. Its selected shallow variants remain limited on precise-recall tasks.

  • The evaluation is limited to autoregressive language modeling under the reported training budgets and model scales.
  • Modular TTT does not exhaust possible inner learners or optimization schedules, including alternatives for deeper and gated learners.Their behavior may differ with other update rules, chunk sizes, optimizers, or hybrid attention designs.
  • The framework currently excludes LaCT-specific Muon-style update normalization and momentum variants because they require specialized fused update paths.LaCT-style checks found no clear quality gain from these refinements.
  • Selected shallow Modular TTT variants remain weaker than LLaMA on containment-style tasks and explicit long-context retrieval.This indicates limitations for precise recall with fixed-state TTT.
  • Table 32 reports family-level RULER accuracies in percent for three multi-NIAH families at 410M and 1.45B, with higher values preferred.
Loading 2608.07110v1…