T2MLR_362M_lstart9_lend24_50B_FineWebEdu
T2MLR (Transformer with Temporal Middle-Layer Recurrence) applies a recurrent connection across a contiguous band of middle layers, carrying a recurrent state between token positions. This checkpoint is a 362M model pretrained from scratch on FineWeb-Edu for ~51.5B tokens.
Configuration
| Parameters | 366.4M |
| Layers | 32 |
| Hidden size | 960 |
| Recurrent band | T2MLR(9,24) — layers 9–24 inclusive, 1-indexed |
l_start / l_end in config.json | 8 / 23 (0-indexed) |
| Recurrent layers | 16 of 32 |
| Mixing module | gated |
| Precision | bfloat16 |
| Training tokens | ~51.5B (100,000 steps) |
| Final training loss | 2.5055 |
Architecture follows SmolLM2 (Llama-style, GQA, SwiGLU, RoPE) with the T2MLR wrapper applied over the middle-layer band. Tokenizer is the SmolLM2 tokenizer (49,152 tokens).
Usage
This model uses a custom wrapper, so load it with the reference implementation rather
than a bare AutoModel call:
git clone https://github.com/princeton-pli/T2MLR.git
cd T2MLR && pip install -r requirements.txt
The loader reads from a local directory, so download the repo first with
snapshot_download rather than passing the repo id straight to it.
import sys, torch
sys.path.insert(0, "T2MLR/src")
from huggingface_hub import snapshot_download
from t2mlr_wrapper import T2MLRWrapper
from transformers import AutoTokenizer
path = snapshot_download("JupiterZhu/T2MLR_362M_lstart9_lend24_50B_FineWebEdu")
model = T2MLRWrapper.from_pretrained_with_t2mlr(path, attn_impl="sdpa", dtype=torch.bfloat16).eval()
tok = AutoTokenizer.from_pretrained(path)
inputs = tok("The capital of France is", return_tensors="pt")
out = model.generate(**inputs, max_new_tokens=32, do_sample=False,
pad_token_id=tok.eos_token_id)
print(tok.decode(out[0], skip_special_tokens=True))
control_flows is required for direct forward calls
When T2MLR is enabled, forward() requires a control_flows tensor shaped like
input_ids. Values <= 1 run the plain (non-recurrent) path; values > 1 mark
positions that participate in recurrence. generate() sets this up for you.
ids = tok("The capital of France is", return_tensors="pt").input_ids
cf = torch.full_like(ids, 2) # 2 => recurrent
logits = model(input_ids=ids, attention_mask=torch.ones_like(ids),
control_flows=cf).logits
Running with control_flows = 1 everywhere disables the recurrence and gives
substantially worse loss — the recurrent band carries a large share of the model's
capability.
Notes
- These are base models trained on a general web corpus with no instruction tuning or alignment. Greedy decoding from short prompts is often repetitive; this is normal for models at this scale and token budget.
- Outputs may be inaccurate, biased, or offensive, reflecting the pretraining data.
Citation
Official implementation: https://github.com/princeton-pli/T2MLR