TRM-text-1B
TRM-text シリーズの 1B パラメータ版。recurrent-depth transformer（RMSNorm + SwiGLU + RoPE + gated residual の TRMBlock を n_layers 層ぶん積み、それを recurrence_steps 回ループさせるコア）を実データ（FineWeb-Edu / FineWeb-2 日本語）で学習。
Status: 実験中のチェックポイント。本格事前学習（Chinchilla目安のトークン数）には未到達。アーキテクチャと学習スクリプトの検証目的のリリース。
アーキテクチャ概要
input_ids
└─ token_emb (vocab × dim)
└─ [ for step in range(recurrence_steps): # 再帰ループ
for block in blocks (n_layers個): # ユニークブロック
x = x + σ(attn_gate) * Attn(RMSNorm(x)) # RoPE付きMHA
x = x + σ(mlp_gate) * SwiGLU(RMSNorm(x))
]
└─ RMSNorm
└─ lm_head (dim × vocab, untied)
実効深度 = n_layers × recurrence_steps。同じ重みを複数回通すことで、パラメータ数を増やさずに test-time compute を増やす（recurrent-depth / "looped transformer" 系の発想。Huginn-3.5B (Geiping et al.), Ouro, AdaPonderLM などと同系統）。
モデル仕様
| 項目 | 値 |
|---|---|
| 総パラメータ数 | ~1.03B |
| 非埋め込みパラメータ | ~411M |
| 埋め込み比 | ~60%（vocabサイズ依存、後述） |
dim | 2048 |
n_layers（ユニークブロック数） | 8 |
recurrence_steps（再帰回数） | 4（推論時に変更可） |
| 実効深度 | 32 |
n_heads / head_dim | 16 / 128 |
| MLP hidden | 5632（SwiGLU） |
| 位置エンコーディング | RoPE |
max_seq_len | 2048 |
| 埋め込み | untied（lm_head と token_emb は別重み） |
| トークナイザ | Qwen2.5 tokenizer 前提（vocab ≈152k）。別トークナイザに変える場合は埋め込み比が大きく変わる |
埋め込み比が高め（vocab≈152kのため）。
recurrence_stepsを増やすほど非埋め込みの実効計算量だけが増える設計。
使い方
from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "summerMC/TRM-text-1B"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True).cuda()
ids = tok("The recurrent-depth transformer", return_tensors="pt").input_ids.cuda()
out = model.generate(ids, max_new_tokens=60, do_sample=True, temperature=0.8, top_p=0.9)
print(tok.decode(out[0], skip_special_tokens=True))
推論時に再帰回数を変える（test-time compute）
学習時より深く"考えさせる"ことができる（劣化する場合もあるので要検証）：
model.config.recurrence_steps = 8 # デフォルト4から増やす
学習
- データ: FineWeb-Edu（英語）70% + FineWeb-2
jpn_Jpan（日本語）30%、ストリーミングpacking - Optimizer: AdamW8bit, lr 3e-4 → 3e-5 cosine, warmup 200 step
- 精度: bf16 autocast、勾配チェックポイント（再帰ループの各ブロック適用ごと）
- ハードウェア: Colab L4 (24GB) — パイプライン検証・短期CPT規模。本格事前学習には未到達
学習・推論コード一式: build_trm1b.py, trm_text_1b_colab.ipynb, configuration_trm_text_ism.py, modeling_trm_text_ism.py
既知の制約 / 注意点
- トークン予算不足: 1Bクラスのtoken-optimal学習（~20Bトークン目安）には到達していない。現状の出力は事前学習不足の挙動を含む可能性が高い。
- Attractor collapse:
recurrence_stepsを学習時より大きく上げすぎると出力が崩れることがある（TRM系列で既知の失敗モード）。recurrence_steps=4を安定運用のデフォルトとして推奨。 - vocab依存の埋め込み比: Qwen2.5 tokenizer (vocab≈152k) 前提。byte-levelなど小vocabに変えると非埋め込み計算の比率が相対的に上がる。
- 評価（HellaSwag/ARC/GSM8K等）は未実施。性能を主張する数値ではない。
ライセンス
MIT。学習データ（FineWeb-Edu / FineWeb-2）のライセンス・利用条件は各データセット側に従う。
Citation
@misc{trm-text-1b,
author = {summerMC},
title = {TRM-text-1B: a recurrent-depth transformer for text},
year = {2026},
url = {https://huggingface.co/summerMC/TRM-text-1B}
}