Tiny Recursive Model (TRM)を試しに走らすだけのメモ

2025年10月20日2025年11月11日

ソース

Oct. 6th 2025にarXivに投稿された次のpaper。

Less is More: Recursive Reasoning with Tiny Networks
Alexia Jolicoeur-Martineau
https://arxiv.org/abs/2510.04871

ソースコードが公開されていますので、試しました。
https://github.com/SamsungSAILMontreal/TinyRecursiveModels

インストール

うちの環境は、intel Z790, core i9-i14900k, rtx5090+rtx4090のローカルパソコン ubuntu gnu/linuxです。

gitコマンドでダウンロードして、venv環境をつくり、そこに入ります。

$ git clone https://github.com/SamsungSAILMontreal/TinyRecursiveModels.git
$ cd TinyRecursiveModels
$ python3 -m venv ./venv
$ source ./venv/bin/activate

あとはreadmeの通り requirementsをいれます。

$ pip install -r requirements.txt # install requirements
$ pip install --no-cache-dir --no-build-isolation adam-atan2

データセットをダウンロードします。今回はテストなので一番軽いSudoku-Extremeで。

$ python dataset/build_sudoku_dataset.py --output-dir data/sudoku-extreme-1k-aug-1000  --subsample-size 1000 --num-aug 1000  # 1000 examples, 1000 augments

トラブルシュート

次のコマンドで学習を実行しようとしますがエラーがでます。

$ run_name="pretrain_mlp_t_sudoku"
python pretrain.py \
arch=trm \
data_paths="[data/sudoku-extreme-1k-aug-1000]" \
evaluators="[]" \
epochs=50000 eval_interval=5000 \
lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0 \
arch.mlp_t=True arch.pos_encodings=none \
arch.L_layers=2 \
arch.H_cycles=3 arch.L_cycles=6 \
+run_name=${run_name} ema=True

ModuleNotFoundError: No module named 'adam_atan2_backend'

HRMのほうのissuesに解決法がありました。
https://github.com/sapientinc/HRM/issues/25

代わりにadam-atan2-pytorchを使うようです。

$ pip install adam-atan2-pytorch
$ sed -i 's/adam_atan2/adam_atan2_pytorch/g' pretrain.py
$ sed -i 's/AdamATan2/AdamAtan2/g' pretrain.py
$ sed -i 's/lr=0,/lr=0.0001,/g' pretrain.py

sedでpretain.pyのadam_atan2を書き換えます。学習率の初期値は0のままだと別のエラーがでました。ので学習率の初期値も適宜設定します。

これでもう一度学習を実行すると走ります。

(もし、fatal error: Python.h: No such file or directoryという別のエラーが出た場合は、sudo apt install python3-devをしてPython.hを入れてください。)

なお、wandbを使っていませんが毎回どうするか聞かれます。

wandb: (1) Create a W&B account
wandb: (2) Use an existing W&B account
wandb: (3) Don't visualize my results
wandb: Enter your choice: 3

次のようにoffline指定しておくと静かになります。

$ export WANDB_MODE=offline

実行時間

L40S GPU 1つでRuntime: < 36 hoursだそうです。

うちだとRTX5090 1つで、6時間くらいかかる予想です。arch.mlp_t=Trueで。

先程の実行コード(pretrain.pyに渡したパラメータ)では50000 epochの学習で5000 epochごとにevalが入ります。

VRAMは 32 GBギリギリまで使っていて、負荷もフル575Wです。gpuは効率よく利用できているようです。

mlpが入らないほうだとVRAMは21 GB使用して、4h30min程度で走りそうです。4.00 it/s –

21 GBだとRTX4090でもそのまま走ります(mlp入りはout of memoryになった)。こちらもやってみると、2.76it/sで6h30minかかるということです。RTX5090のほうが44%速いです。

RTX5090+RTX4090両方を使うと、5.25 it/sでした。RTX5090の速度からすると2枚使う割に1.3倍にしかなっておらず遅い気がしますが、RTX4090で律速すると考えると5.25/2.76 = 1.9なので、ほぼ4090 2枚分の速度ということで悪くはないです。RTX4090が遅いということで…(1~2年前までは爆速のイメージでしたが)。

結局evalも時間がかかってmlpが入らないほうで8時間10分かかりました。

wandb: Run history:
wandb:            num_params ▁
wandb:        train/accuracy ▁▁▁▆▆▆▇▇▇▇▇█▇█▇▇▇██▇████████▇█████▇█████
wandb:           train/count ▁▁▁▁████████████████████████████████████
wandb:  train/exact_accuracy ▁▁▁▁▁▄▂▂▄▅▅▅▅▅▅▅▅▅▅▆█▆█▇▇▇█▇█▇▆▇▇▇▇█████
wandb:         train/lm_loss █▄▂▂▂▂▂▂▂▂▂▁▂▂▂▁▂▁▁▁▁▂▁▁▁▁▁▁▁▁▂▁▁▁▁▁▁▁▁▁
wandb:              train/lr ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
wandb: train/q_halt_accuracy ▁▁█████████████████████▇████████████████
wandb:     train/q_halt_loss ▃▁▁▁▃▇▂▂█▂▂▂▁▂▁▂▁▁▂▅▂▆▁▆▄▁▃▁▂▄▁▂▂▁▂▂▂▃▂▅
wandb:           train/steps ▁██▇▇▆▇▆▇▆▇▇▆▆▆▆▆▅▆▆▅▅▆▅▅▅▄▄▅▆▅▅▄▅▅▅▅█▅▅
wandb: 
wandb: Run summary:
wandb:            num_params 6828034
wandb:        train/accuracy 0.85783
wandb:           train/count 1
wandb:  train/exact_accuracy 0.55914
wandb:         train/lm_loss 0.66863
wandb:              train/lr 0.0001
wandb: train/q_halt_accuracy 0.98925
wandb:     train/q_halt_loss 0.01033
wandb:           train/steps 8.6129

65100/65104 [8:10:37<00:01,  2.21it/s]

もっと短い時間で試したいとき

実行パラメータは、config/cfg_pretrain.yamlにかかれているものがデフォルト値です(パラメータ管理にはhydraが使われています)。

実行時に、pretrain.py global_batch_size=32 epochs=1 eval_interval=1などとすると、パラメータを上書きできます。

使用するモデルは archで渡します。arch=trmだとTRMを使い、arch=hrmだとHRMを使います。モデルの構造はconfig/arch/trm.yamlなどに書かれています。

ARC-AGI-1を試す

データセットを作成します。

# ARC-AGI-1
python -m dataset.build_arc_dataset \
  --input-file-prefix kaggle/combined/arc-agi \
  --output-dir data/arc1concept-aug-1000 \
  --subsets training evaluation concept \
  --test-set-name evaluation

たとえば次のように実行します。

CUDA_VISIBLE_DEVICES=0 python pretrain.py \
arch=trm \
data_paths="[data/arc1concept-aug-1000]" \
arch.L_layers=2 \
arch.H_cycles=3 arch.L_cycles=4 \
ema=True

でもかなり時間がかかりそうですので、はじめは軽量にしてテストするのがよいです。

CUDA_VISIBLE_DEVICES=0 python pretrain.py arch=trm_singlez data_paths=["data/arc1concept-aug-1000"] global_batch_size=32 epochs=1 eval_interval=1

リポジトリの概要 (生成した解説コピペ)

以下では、pretrain.py を中心に、TinyRecursiveModels リポジトリの主要なコードの処理内容をまとめます。

● このコードは一言で言うと何をしているか

TinyRecursiveModels の各種モデル（TRM や階層型モデルなど）を，パズルデータセット上で事前学習（pretrain）し，ログ出力やチェックポイント保存を行うトレーニングスクリプトです。

● 主要な処理ステップ

設定の読み込み
- config/cfg_pretrain.yaml からハイパーパラメータやモデルアーキテクチャを読み込む。
データ準備
- puzzle_dataset.py の PuzzleDataset でパズル問題データをロードし、DataLoader を生成。
モデル・最適化器の初期化
- models/recursive_reasoning 以下のサブモジュールから指定モデルを構築。
- 学習率や重み減衰などを設定し，複数のオプティマイザを用意。
トレーニングループ
- 各エポック・バッチでフォワード／バックワードを実行し，損失を計算，パラメータ更新。
- wandb でメトリクスログを送りつつ，定期的にモデルを保存。
オプショナルな重みミックス
- mix_weights_direct 関数で複数モデル間の重みをα比率で合成し，新たなモデルを生成。

● 主要なクラス／関数と役割

名前	役割
`PuzzleDataset` （`puzzle_dataset.py`）	パズル問題を読み込み，入力（シーケンス／特徴量）とラベルを返す Dataset
`build_model()` （`pretrain.py` 内）	設定に応じて TRM／HRM／Transformer ベースラインなどを構築
`mix_weights_direct()`	複数モデルの state_dict を α 合成し，新モデルに適用
`train_one_epoch()`	1 エポック分のフォワード・バックワード・更新・ログを実行
`main()`	全体の処理を統括：設定読み込み→データ準備→モデル初期化→ループ

● データの入出力フロー

入力
- 形式・型：テキスト／数値シーケンス（例：長さ N の整数トークン列）
- 具体例：[12, 5, 27, 9, …] と正解ラベル 3
処理
- 埋め込み層でトークンをベクトル化 → 再帰的／階層的モデルに入力 → 最終層で分類／回帰出力
- 損失計算（クロスエントロピーなど）→ 勾配降下
出力
- 型：ログ（wandb）、モデルチェックポイント（.pt）
- 具体例：epoch=10, loss=0.1234 のログと，model_epoch10.pt

● pretrain.pyのフローチャート

Core 300とCore 200Sの違い比較まとめ【Levoit 空気清浄機】