Modular TTT: Rethinking Test-Time Training as Composable Modules

2026-08-07

AI Summary

This work addresses the inflexibility of existing test-time training (TTT) methods, which are typically implemented as monolithic, hard-coded systems that hinder modular design and component-wise analysis. To overcome this limitation, the authors propose the first modular TTT framework, modeling the internal learner as a directed acyclic graph that explicitly decouples key elements such as fast weight networks, loss functions, and learning rates. The framework automatically composes elementary forward, backward, and query rules to construct complete computational pipelines, enabling systematic ablation studies and flexible reconfiguration. Through this approach, the study reveals the critical roles of small learning rate initialization, weight decay, and single-layer nonlinearity in achieving strong performance. Models built within this framework—scaled to 410 million and 1.45 billion parameters and trained on 100 billion tokens—match the training loss and downstream performance of Gated DeltaNet.

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.

Problem

• Test-Time Training
• modular design
• fast weights
• sequence modeling
• component analysis

Innovation

• Modular TTT
• fast weights
• test-time training
• composable modules
• directed acyclic graph