👋 Hi, everyone!
We are ByteDance Seed team.

You can get to know us better through the following channels👇

Modular TTT

This repository contains the research code for Modular TTT: Rethinking Test-Time Training as Composable Modules.

Modular TTT treats test-time training (TTT) layers as composable inner-learning graphs rather than as separately hard-coded model variants. A TTT memory is represented as a directed acyclic graph whose nodes are primitive operations such as Linear, Act, Norm, Add, Mul, and loss functions. Given local train-view, train-view backward, and query-view rules for each primitive, the framework composes the full TTT computation automatically.

This code release includes the Modular TTT package, custom operators, model configs, training scripts, evaluation scripts, and the LLaMA/Gated DeltaNet baselines used for comparison.

Main Features

• Graph-structured TTT memory modules built from composable primitives.
• Local train-view, backward-transport, and query-view rules for TTT updates.
• Configurable loss functions, inner learning rates, decay, normalization, and activation choices.
• Hugging Face-style TTTConfig, TTTModel, and TTTForCausalLM interfaces.
• Custom operators for efficient fast-weight update and query-view computation.
• Flame-based training, checkpoint conversion, and lm-eval-harness evaluation scripts.
• RULER and per-token loss evaluation entry points.

Repository Layout

modular_ttt/              Core Modular TTT Python package.
xopes/                    Custom operators used by Modular TTT and xmixers.
xmixers/                  Token mixer implementations and baseline modules.
flame/                    Training framework, model configs, and trainer code.
flame_script/env/         Environment setup script for this repository.
flame_script/flame_task/  Training, checkpoint conversion, and lm-eval scripts.
flame_script/ruler/       RULER evaluation scripts.
flame_script/per_token_loss/
                           Notes for per-token loss evaluation.
debug_ttt.sh              Paper-related training and evaluation commands.

Installation

Use a Linux environment with a working PyTorch installation. From the repository root:

git clone https://github.com/ByteDance-Seed/Modular-TTT.git
cd Modular-TTT

bash flame_script/env/setup_env.sh

The setup script installs the local packages in this repository and the external Python dependencies used by the training and evaluation scripts, including flash-linear-attention, bitsandbytes, lm_eval, datasets, wandb, torchtitan, and torchdata.

If you run the setup script from a different working directory, set repo_dir explicitly:

repo_dir=/path/to/modular-ttt bash /path/to/modular-ttt/flame_script/env/setup_env.sh

Data Setup

Training data is not included in this repository. Set the data paths before launching training:

export DATA_10B_PATH=/path/to/your/10b_training_data
export DATA_100B_PATH=/path/to/your/100b_training_data

The scripts pass these paths to the Flame data loader. The optional DATASET_NAME variable defaults to default:

export DATASET_NAME=default

By default, experiment outputs, checkpoints, and caches are written under the repository directory. To keep outputs elsewhere:

export home_dir=/path/to/your/workdir

Quick Start

The main training entry point is flame_script/flame_task/script.sh. It trains a model, saves checkpoints, converts the target checkpoint to Hugging Face format, and runs lm-eval-harness.

The following command runs a short end-to-end workflow with 2 training steps and wikitext evaluation:

cd /path/to/modular-ttt
export DATA_10B_PATH=/path/to/your/10b_training_data

config_dir=modular_ttt/ttt_linear_graph_mse_no_norm_sd_sinlr_official_init

bash flame_script/flame_task/script.sh \
  1 false \
  ${config_dir} \
  ttt_linear_graph_mse_no_norm_sd_sinlr_90m_official_init \
  2 1 128 false \
  ttt_linear_mse_no_norm_sd \
  2 1 3e-4 1 10b -1 none wikitext

Use the paper-style commands in debug_ttt.sh for the full experiment settings.

Reproducing Paper-Style Runs

debug_ttt.sh collects the paper-related commands for:

• Modular TTT scale-up variants.
• Modular TTT loss, decay, nonlinearity, and deep-memory ablations.
• LLaMA baseline runs.
• Gated DeltaNet baseline runs.

Set the data paths first, then open debug_ttt.sh and uncomment the command you want to run:

cd /path/to/modular-ttt
export DATA_10B_PATH=/path/to/your/10b_training_data
export DATA_100B_PATH=/path/to/your/100b_training_data

bash debug_ttt.sh

The script uses historical config names such as 90m, 310m, and 1_2b. These names refer roughly to the non-vocabulary backbone size: transformer blocks, token mixer, channel mixer, and norms. The paper reports total parameter counts, including token embeddings and the untied language-model head. Therefore the mapping used in the paper is:

Config name	Paper scale
90m	160M
310m	410M
1_2b	1.45B

Evaluation

The default flame_script/flame_task/script.sh path evaluates the converted checkpoint with lm-eval-harness. By default it covers:

• Perplexity: wikitext, lambada_openai
• Multiple choice: boolq, piqa, social_iqa, hellaswag, winogrande, arc_easy, arc_challenge, openbookqa
• Containment-style tasks: fda, swde, squad_completion

You can override the task list with the final argument to script.sh, for example:

bash flame_script/flame_task/script.sh ... wikitext,lambada_openai

RULER commands are listed in:

flame_script/ruler/script.sh

The minimal per-token loss evaluator is in:

flame_script/per_token_loss/readme.md

Core Package Usage

After installing the repository dependencies:

from modular_ttt import TTTConfig, TTTForCausalLM

config = TTTConfig()
model = TTTForCausalLM(config)

For the standalone package details, see:

modular_ttt/README.md

Notes

• This repository releases code and configs, not training data or model checkpoints.
• The primary implementation is modular_ttt, but efficient execution depends on custom operators from xopes.
• flame and xmixers are included to make the training and evaluation commands reproducible from this repository.
• The paper focuses on autoregressive language modeling. The current shallow Modular TTT variants are competitive with recurrent baselines in several settings, but exact long-context retrieval remains an open limitation.

Citation

If you find this work useful, please cite our paper:

@misc{tang2026modulartttrethinkingtesttime,
  title={Modular TTT: Rethinking Test-Time Training as Composable Modules},
  author={Bohao Tang and Zhen Qin and Yuqi Pan and Zheng Li and Pengfei Liu and Ya Zhang},
  year={2026},
  eprint={2608.07110},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2608.07110},
}

Security

If you discover a potential security issue in this project, please notify ByteDance Security through the security center or vulnerability reporting email.

Please do not create a public GitHub issue for security vulnerabilities.

License

This project is licensed under the Apache License 2.0.

This repository bundles third-party subcomponents with their own licenses. See THIRD_PARTY_NOTICES.md for third-party copyright and license notices.

About ByteDance Seed Team

Founded in 2023, ByteDance Seed Team is dedicated to crafting the industry's most advanced AI foundation models. The team aspires to become a world-class research team and make significant contributions to the advancement of science and society.