Tiny Recursive Models (TRM)

repository·main·Indexed 27 days ago

https://github.com/samsungsailmontreal/tinyrecursivemodels

A codebase for recursive reasoning using small neural networks (e.g., 7M parameters). TRM enables models to progressively improve answers through recursive latent and answer updates to solve complex reasoning tasks such as ARC-AGI, Sudoku-Extreme, and Maze-Hard.

Tokens: 1.9K
Snippets: 6
Records: 6
Agent score: 42%

What's inside Tiny Recursive Models

Prepare Sudoku-Extreme and Maze-Hard datasets
main
Generate datasets for Sudoku-Extreme and Maze-Hard using the provided build scripts.
```
# Sudoku-Extreme
python dataset/build_sudoku_dataset.py --output-dir data/sudoku-extreme-1k-aug-1000  --subsample-size 1000 --num-aug 1000

# Maze-Hard
python dataset/build_maze_dataset.py
```
https://github.com/samsungsailmontreal/tinyrecursivemodels/blob/main/README.md
Install Tiny Recursive Models (TRM)
main
To set up the TRM environment, ensure you have Python 3.10 (or similar) and CUDA 12.6.0 (or similar) installed. Follow these steps to install the necessary dependencies:
1. Upgrade core build tools.
2. Install the nightly PyTorch build corresponding to your CUDA version.
3. Install project requirements.
4. Install adam-atan2 without cache or build isolation.
5. (Optional) Log in to Weights & Biases to sync results.
```
pip install --upgrade pip wheel setuptools
pip install --pre --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu126
pip install -r requirements.txt
pip install --no-cache-dir --no-build-isolation adam-atan2 
wandb login YOUR-LOGIN
```
https://github.com/samsungsailmontreal/tinyrecursivemodels/blob/main/README.md

Prepare ARC-AGI-1 and ARC-AGI-2 datasets

main

Use the dataset.build_arc_dataset module to prepare datasets for ARC-AGI.

Warning: You cannot train on both ARC-AGI-1 and ARC-AGI-2 and evaluate them both because ARC-AGI-2 training data contains some ARC-AGI-1 evaluation data.

# 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

# ARC-AGI-2
python -m dataset.build_arc_dataset \
  --input-file-prefix kaggle/combined/arc-agi \
  --output-dir data/arc2concept-aug-1000 \
  --subsets training2 evaluation2 concept \
  --test-set-name evaluation2

Run ARC-AGI-1 and ARC-AGI-2 experiments

main

Train a TRM model on ARC-AGI datasets using distributed training (4 H-100 GPUs). Runtime is approximately 3 days per experiment.

# ARC-AGI-1
run_name="pretrain_att_arc1concept_4"
torchrun --nproc-per-node 4 --rdzv_backend=c10d --rdzv_endpoint=localhost:0 --nnodes=1 pretrain.py \
arch=trm \
data_paths="[data/arc1concept-aug-1000]" \
arch.L_layers=2 \
arch.H_cycles=3 arch.L_cycles=4 \
+run_name=${run_name} ema=True

# ARC-AGI-2
run_name="pretrain_att_arc2concept_4"
torchrun --nproc-per-node 4 --rdzv_backend=c10d --rdzv_endpoint=localhost:0 --nnodes=1 pretrain.py \
arch=trm \
data_paths="[data/arc2concept-aug-1000]" \
arch.L_layers=2 \
arch.H_cycles=3 arch.L_cycles=4 \
+run_name=${run_name} ema=True

Run Sudoku-Extreme experiments

main

Train a TRM model on the Sudoku-Extreme dataset. You can choose between an MLP-based architecture or an attention-based architecture.

Hardware Note: Training on 1 L40S GPU takes approximately 18 hours.

# MLP-based architecture (Expected ~87% accuracy)
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

# Attention-based architecture (Expected ~75% accuracy)
run_name="pretrain_att_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.L_layers=2 \
arch.H_cycles=3 arch.L_cycles=6 \
+run_name=${run_name} ema=True

Run Maze-Hard experiments

main

Train a TRM model on the Maze-Hard dataset. You can run this using distributed training (4 L40S GPUs) or on a single L40S GPU by reducing the global_batch_size.

# Distributed training (4 L40S GPUs)
run_name="pretrain_att_maze30x30"
torchrun --nproc-per-node 4 --rdzv_backend=c10d --rdzv_endpoint=localhost:0 --nnodes=1 pretrain.py \
arch=trm \
data_paths="[data/maze-30x30-hard-1k]" \
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.L_layers=2 \
arch.H_cycles=3 arch.L_cycles=4 \
+run_name=${run_name} ema=True

# Single GPU training (1 L40S GPU)
run_name="pretrain_att_maze30x30_1gpu"
python pretrain.py \
arch=trm \
data_paths="[data/maze-30x30-hard-1k]" \
evaluators="[]" \
epochs=50000 eval_interval=5000 \
lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0 global_batch_size=128 \
arch.L_layers=2 \
arch.H_cycles=3 arch.L_cycles=4 \
+run_name=${run_name} ema=True