T2MLR_362M_lstart9_lend24_50B_FineWebEdu

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

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