TRM-LM-v2-10M

TRM-LM v2 10M

A 9.8M-parameter language model that thinks in loops, by SmallAICreator / UltraLabs.

It is an ordinary 7-layer transformer whose middle 5 layers are run several times per token (prelude 1 → core 5 × r loops → coda 1). Same parameters, more thinking. With r = 1 it is a normal 7-layer model; every extra loop re-reads the context with the same weights.

The result: same size, same data, better

Four 10M models trained on the exact same 602,144,768 tokens in the same order (FineWeb-Edu, Cosmopedia-v2, SlimPajama), same optimizer and schedule, same 16K tokenizer. Only the architecture differs.

LAMBADA +38% over a normal model of the same size. For scale: Pythia-70M (7× the parameters, ~500× the data) scores 22.7%.

Thinking longer helps (same weights, loop count chosen at run time)

It gains where the answer depends on using the context (LAMBADA, SciQ, perplexity) and ties where it depends on stored knowledge (HellaSwag, PIQA). Loops add computation, not memory. The cost: at 3 loops it runs ~17 layer passes per token versus 7, about 2.4× the compute of the normal model.

Files

import torch
from tokenizers import Tokenizer
from modeling_trm2 import TRM2, generate

tok = Tokenizer.from_file("tokenizer.json")
model = TRM2()
model.load_state_dict(torch.load("trm2_sft.pt", map_location="cpu"))
model.eval()
print(generate(model, tok, "What is the capital of France?", loops=3))
# -> "The capital of France is Paris, ..."

Chat format: <|user|>message<|end|><|assistant|>reply<|end|> (ids 3, 5, 4, 5). generate(..., chat=False) continues plain text with the base model.

Chat model: be realistic

It is 10M parameters. It greets you ("hi" → "Hello! How can I help you today?"), answers simple factual questions ("The capital of France is Paris."), and says "You're welcome!" to thanks. It still drifts on jokes, poems, arithmetic and questions about itself. Chat tuning also costs some of the base model's edge: LAMBADA drops from 18.7% (base) to 13.1% at 3 loops / 13.4% at 4 loops. Use trm2.pt for text continuation and benchmarks, trm2_sft.pt for chat. Held-out chat loss is lowest at 3 loops (1: 2.305, 2: 2.143, 3: 2.118, 4: 2.123, 6: 2.175).

Training

• Architecture: D=256, 4 heads (dim 64), SwiGLU 704, RoPE, pre-RMSNorm, tied embeddings, context 1024. Each loop i adds a learned gain and tag vector so the shared core can behave differently per pass.
• Pretraining: 4,594 steps × 131,072 tokens = 602M tokens on 2× T4 (Kaggle), AdamW (0.9, 0.95), weight decay 0.1, peak LR 1.5e-3, 400 warmup, cosine to 10%, fp16. Loop count sampled per step from {1, 2, 3, 4} with probabilities {0.1, 0.2, 0.35, 0.35}; the loss is on the final output only and gradients flow through every loop.
• Chat SFT: 40 min, 178M tokens: 150K smol-smoltalk conversations, everyday-conversations ×8, identity and greeting examples; loss on assistant tokens only; LR 5e-4.

Why v1 lost and v2 wins. TRM-LM v1 followed the original Tiny Recursive Model recipe (two shared layers, answer/scratchpad states, detached carry between cycles, every cycle supervised). On text it learned to do all its work in the first cycle, and extra cycles made it slightly worse. v2 moves the parameters into depth, lets gradients cross loops, and supervises only the final output, and extra loops then help.

Made by SmallAICreator / UltraLabs.