zeyun-zhong/e2-ttt-mlp-340M-15B

E²-TTT_MLP 340M

E²-TTT is a test-time-training (TTT) layer whose chunk-wise update exactly reproduces the chunk-end fast-weight and momentum states of the per-token recurrence, under the standard approximation of taking inner-loop gradients at the chunk-start weights. Training stays fully chunk-parallel while the per-token learning-rate, momentum and decay dynamics are preserved.

This checkpoint is the MLP fast-weight variant at the 340M scale, trained on ~15B tokens of FineWeb-Edu. It is one of four released checkpoints:

	340M	1.3B
MLP fast weights	zeyun-zhong/e2-ttt-mlp-340M-15B	zeyun-zhong/e2-ttt-mlp-1.3B-15B
SwiGLU fast weights	zeyun-zhong/e2-ttt-swiglu-340M-15B	zeyun-zhong/e2-ttt-swiglu-1.3B-15B

• 📄 Paper: arXiv:2608.21308
• 💻 Code: github.com/zeyun-zhong/E2-TTT

Usage

The architecture lives in the e2_ttt package, which registers e2_ttt_mlp with transformers. Install it first — trust_remote_code alone is not enough, this repository ships no modeling code:

pip install git+https://github.com/zeyun-zhong/E2-TTT

import torch
import e2_ttt  # registers e2_ttt_mlp / e2_ttt_swiglu
from transformers import AutoModelForCausalLM, AutoTokenizer

name = "zeyun-zhong/e2-ttt-mlp-340M-15B"
tokenizer = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.bfloat16).cuda()

ids = tokenizer("The capital of France is", return_tensors="pt").input_ids.cuda()
print(tokenizer.decode(model.generate(ids, max_new_tokens=16)[0]))

Pass dtype explicitly. The weights are stored in bfloat16, but from_pretrained defaults to fp32 regardless of what config.json says; use dtype=torch.bfloat16 (or dtype="auto") to load the model in the precision it was evaluated in. A CUDA GPU is required — the TTT kernels and the sliding-window attention branch use FlashAttention and Triton.

Model details

Architecture	E²-TTT_MLP (e2_ttt_mlp)
Parameters	424,841,408 total, 359,305,408 non-embedding
Layers / hidden size / heads	24 / 1024 / 8
TTT chunk size C	512
Sliding-window attention window	512
Inner-loop peak LR / weight decay	0.01 / 0.1
Momentum / decay / closed-form transition	enabled
Vocabulary	32,000 (fla-hub/transformer-1.3B-100B tokenizer)
Training sequence length	2048
Precision	bfloat16

Each layer runs a TTT fast-weight branch alongside a sliding-window attention branch, combined by a learned memory gate; the fast weights are updated once per chunk of 512 tokens using the paper's closed-form state transition.

Training

• Data: HuggingFaceFW/fineweb-edu (sample-100BT): 28,672 steps × 256 sequences × 2048 tokens = 15.0B tokens (the paper reports both scales as a 15B-token budget).
• Optimizer: AdamW, peak LR 1e-3, cosine decay to 0.1× with 1,024 warm-up steps, weight decay 0.1, gradient clipping 1.0.
• Precision: bf16 mixed precision (FSDP2 MixedPrecisionPolicy) over fp32 master weights; the released weights are the fp32 master weights cast to bfloat16.
• Framework: flame (torchtitan); see training/train_340M.sh in the repository for the exact command.

Evaluation

Benchmark results (Wikitext / LAMBADA perplexity, commonsense reasoning, SQuAD / SWDE recall, RULER S-NIAH and LongBench) are reported in the paper. To reproduce them, use training/eval.sh from the repository, which drives lm-evaluation-harness at commit c9772b9 with dtype=bfloat16.

These models are trained at a ~2K context, so every longer evaluation (RULER, LongBench) is extrapolation. max_position_embeddings is 2048, but it is not a hard cap: the rotary embedding is local to the 512-token attention window and the TTT branch carries the rest of the context in its fast weights.

Limitations

A 340M model trained on ~15B tokens of web text is a research artifact for studying TTT layers, not a general-purpose assistant. It is not instruction-tuned or aligned, and it will produce inaccurate, biased or otherwise undesirable text.

Citation

@article{zhong2026e2ttt,
  title   = {Rethinking Expressivity and Efficiency in Test-Time Training},
  author  = {Zhong, Zeyun and Chen, Joya and Martin, Manuel and
             Diederichs, Frederik and Gall, Juergen and Beyerer, Juergen},
  journal = {arXiv preprint arXiv:2608.21308},
  year    = {2026}
}