# Modular-TTT

**Repository Path**: ByteDance-Seed/Modular-TTT

## Basic Information

- **Project Name**: Modular-TTT
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Apache-2.0
- **Default Branch**: main
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-09-15
- **Last Updated**: 2026-09-17

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

<div align="center">
  👋 Hi, everyone!
  <br>
  We are <b>ByteDance Seed team.</b>
</div>
<p align="center">
  You can get to know us better through the following channels👇
  <br>
  <a href="https://seed.bytedance.com/">
    <img src="https://img.shields.io/badge/Website-%231e37ff?style=for-the-badge&logo=bytedance&logoColor=white"></a>
  <a href="https://github.com/user-attachments/assets/5793e67c-79bb-4a59-811a-fcc7ed510bd4">
    <img src="https://img.shields.io/badge/WeChat-07C160?style=for-the-badge&logo=wechat&logoColor=white"></a>
  <a href="https://www.xiaohongshu.com/user/profile/668e7e15000000000303157d?xsec_token=ABl2-aqekpytY6A8TuxjrwnZskU-6BsMRE_ufQQaSAvjc%3D&xsec_source=pc_search">
    <img src="https://img.shields.io/badge/Xiaohongshu-%23FF2442?style=for-the-badge&logo=xiaohongshu&logoColor=white"></a>
  <a href="https://www.zhihu.com/org/dou-bao-da-mo-xing-tuan-dui/">
    <img src="https://img.shields.io/badge/zhihu-%230084FF?style=for-the-badge&logo=zhihu&logoColor=white"></a>
</p>

![seed logo](https://github.com/user-attachments/assets/c42e675e-497c-4508-8bb9-093ad4d1f216)

# [Modular TTT](https://arxiv.org/abs/2608.07110)

<p align="center">
  <a href="https://arxiv.org/abs/2608.07110">
    <img src="https://img.shields.io/badge/arXiv-2608.07110-b31b1b?style=for-the-badge&logo=arxiv&logoColor=white" alt="arXiv"></a>
</p>

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.

![Modular TTT framework overview](assets/framework.png)

## 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

```text
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:

```bash
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:

```bash
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:

```bash
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`:

```bash
export DATASET_NAME=default
```

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

```bash
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:

```bash
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:

```bash
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
bash flame_script/flame_task/script.sh ... wikitext,lambada_openai
```

RULER commands are listed in:

```text
flame_script/ruler/script.sh
```

The minimal per-token loss evaluator is in:

```text
flame_script/per_token_loss/readme.md
```

## Core Package Usage

After installing the repository dependencies:

```python
from modular_ttt import TTTConfig, TTTForCausalLM

config = TTTConfig()
model = TTTForCausalLM(config)
```

For the standalone package details, see:

```text
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](https://arxiv.org/abs/2608.07110):

```bibtex
@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](https://security.bytedance.com/src)
or [vulnerability reporting email](mailto:src@bytedance.com).

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

## License

This project is licensed under the [Apache License 2.0](LICENSE).

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

## About [ByteDance Seed Team](https://seed.bytedance.com/)

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.