ARC-AGI: Recursive and Hierarchical Reasoning Models

This repository contains PyTorch implementations of two state-of-the-art architectures designed for the ARC-AGI (Abstraction and Reasoning Corpus) benchmark:

1. Tiny Recursive Model (TRM): A parameter-efficient model using nested recursive loops to refine reasoning ("Thinking") and answers ("Answering").
2. Hierarchical Reasoning Model (HRM): A dual-module architecture (Fast $L$-module, Slow $H$-module) inspired by System 1 / System 2 cognitive theories.

Project Structure

├── data/
│   ├── arc-agi_training_challenges.json
│   └── ...
├── models/
│   ├── trm.py          # Tiny Recursive Model architecture
│   └── hrm.py          # Hierarchical Reasoning Model architecture
├── dataset.py          # Data loading and augmentation (flip/rotate)
├── train_trm.py        # Training script for TRM (Deep Supervision)
├── train_hrm.py        # Training script for HRM
├── inference_new.py    # Unified inference script
├── test_trm.py         # Unit tests for TRM
└── test_hrm.py         # Unit tests for HRM

Models

Tiny Recursive Model (TRM)

• Concept: Uses a single shared Transformer block recursively.
• Mechanism:
  • Inner Loop: Updates latent state $z$ (Thinking) while keeping answer hypothesis $y$ fixed.
  • Outer Loop: Updates answer hypothesis $y$ using refined $z$.
• Key File: models/trm.py

Hierarchical Reasoning Model (HRM)

• Concept: Separate modules for perception and planning.
• Mechanism:
  • L-Module (Fast): Updates frequently, solving local tactics.
  • H-Module (Slow): Updates sparsely, providing global context/planning.
• Key File: models/hrm.py

Usage

1. Installation

Ensure you have PyTorch installed.

pip install torch tqdm numpy

2. Data

Place your ARC-AGI JSON data files in the data/ directory. Expected standard format: arc-agi_training_challenges.json.

3. Training

Train TRM:

python train_trm.py

Note: TRM uses heavy recursion. If you encounter OOM errors (especially on MPS/Mac), reduce max_seq_len in train_trm.py or batch size.

Train HRM:

python train_hrm.py

4. Inference

To run inference on a sample task using a trained model:

python inference_new.py

Edit inference_new.py to point to your specific checkpoint path.

Credits

Based on the 2025 analysis of ARC-AGI architectures: "The Renaissance of Algorithmic Reasoning", implementing concepts from Samsung SAIL Montréal (TRM) and Sapient Intelligence (HRM).