Views
No views yet
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を増やすほど非埋め込みの実効計算量だけが増える設計。
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3repo = "summerMC/TRM-text-1B"
4tok = AutoTokenizer.from_pretrained(repo)
5model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True).cuda()
6
7ids = tok("The recurrent-depth transformer", return_tensors="pt").input_ids.cuda()
8out = model.generate(ids, max_new_tokens=60, do_sample=True, temperature=0.8, top_p=0.9)
9print(tok.decode(out[0], skip_special_tokens=True))model.config.recurrence_steps = 8 # デフォルト4から増やすjpn_Jpan(日本語)30%、ストリーミングpackingbuild_trm1b.py, trm_text_1b_colab.ipynb, configuration_trm_text_ism.py, modeling_trm_text_ism.pyrecurrence_steps を学習時より大きく上げすぎると出力が崩れることがある(TRM系列で既知の失敗モード)。recurrence_steps=4 を安定運用のデフォルトとして推奨。1@misc{trm-text-1b,
2 author = {summerMC},
3 title = {TRM-text-1B: a recurrent-depth transformer for text},
4 year = {2026},
5 url = {https://huggingface.co/summerMC/TRM-text-1B}
6}