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}
}