Modular TTT turns the TTT inner learner into a composable DAG; ablations favor shallow linear memory

Modular TTT: Rethinking Test-Time Training as Composable Modules

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

cs.LG, cs.CL

2026-08-07

Modular TTT casts the TTT inner learner as a DAG of independent dimensions; ablations favor small lr, decay, and one nonlinearity, and its 1.45B model matches Gated DeltaNet.

What problem this solves

Test-time training (TTT) frames sequence modeling as online learning: as the model processes a sequence, an internal "fast weight" is updated by a learning rule. TTT variants are proliferating, but each is hard-coded. Changing one component (fast-weight network, loss, learning rate, decay, normalization) often changes several at once, which makes new variants hard to design and obscures what each component actually contributes.

Method

Modular TTT represents the inner learner as a directed acyclic graph (DAG): nodes are registered primitives (linear layer, nonlinearity, normalization, loss) and edges are tensor dependencies. The framework automatically composes the train-view forward and backward and the causal query-view rules into the full graph-level TTT computation, including the fast-weight state transition. Swapping a component means swapping a node or subgraph, so each dimension can be varied in isolation. From ByteDance Seed and SJTU.

Results

The systematic ablations (at 160M and 410M scale) give a few clear conclusions:

Scaling the selected variant to 410M and 1.45B on 100B tokens: at 1.45B, training loss 2.3150 and multiple-choice average 56.4%, on par with Gated DeltaNet (2.3042 / 56.3%) and close to LLaMA (2.3041). Training throughput is 2.2x to 3.3x the official TTT implementation.

Why it matters

There are two layers of value. In engineering, the DAG plus automatic composition turns "build a new TTT variant and compare it fairly" from hand-crafting into assembling blocks, and the analytic backward operators bring a 2 to 3x throughput gain. In understanding, it maps the TTT design space more clearly: go shallow with one nonlinearity and decay, not deep with normalization. The paper also gives a mathematical explanation for why deep fast weights are hard to optimize, namely that different factorizations of the same effective fast weight induce very different update directions.

Limitations

Scale and context are the weak spots. On the RULER long-context benchmark, LLaMA remains substantially stronger (especially at 8k), so the current shallow Modular TTT still struggles with precise long-context recall; containment-style tasks are also weak. 1.45B on 100B tokens is modest by today's standards, and whether parity holds at larger scale is unverified. "On par with GDN" means matching another linear-attention or recurrent baseline, not beating it.

Terms

Source

What people are saying

Related papers

All paper explainers