README.md

---
license: mit
language:
- en
- ja
library_name: transformers
tags:
- recurrent-depth
- causal-lm
- trm-text
- pytorch
pipeline_tag: text-generation
---

# 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` を増やすほど非埋め込みの実効計算量だけが増える設計。

## 使い方

```python
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）

学習時より深く"考えさせる"ことができる（劣化する場合もあるので要検証）：

```python
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`](./build_trm1b.py), [`trm_text_1b_colab.ipynb`](./trm_text_1b_colab.ipynb), [`configuration_trm_text_ism.py`](./configuration_trm_text_ism.py), [`modeling_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

```bibtex
@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}
}
```