HRM README

The project's README from github.com/sapientinc/HRM. As one markdown file: README.md.

Hierarchical Reasoning Model

Reasoning, the process of devising and executing complex goal-oriented action sequences, remains a critical challenge in AI. Current large language models (LLMs) primarily employ Chain-of-Thought (CoT) techniques, which suffer from brittle task decomposition, extensive data requirements, and high latency. Inspired by the hierarchical and multi-timescale processing in the human brain, we propose the Hierarchical Reasoning Model (HRM), a novel recurrent architecture that attains significant computational depth while maintaining both training stability and efficiency. HRM executes sequential reasoning tasks in a single forward pass without explicit supervision of the intermediate process, through two interdependent recurrent modules: a high-level module responsible for slow, abstract planning, and a low-level module handling rapid, detailed computations. With only 27 million parameters, HRM achieves exceptional performance on complex reasoning tasks using only 1000 training samples. The model operates without pre-training or CoT data, yet achieves nearly perfect performance on challenging tasks including complex Sudoku puzzles and optimal path finding in large mazes. Furthermore, HRM outperforms much larger models with significantly longer context windows on the Abstraction and Reasoning Corpus (ARC), a key benchmark for measuring artificial general intelligence capabilities. These results underscore HRM’s potential as a transformative advancement toward universal computation and general-purpose reasoning systems.

Read Our Paper: https://arxiv.org/abs/2506.21734

Join Our Discord Community: https://discord.gg/sapient

Quick Start Guide 🚀

Prerequisites ⚙️

Ensure PyTorch and CUDA are installed. The repo needs CUDA extensions to be built. If not present, run the following commands:

# Install CUDA 12.6
CUDA_URL=https://developer.download.nvidia.com/compute/cuda/12.6.3/local_installers/cuda_12.6.3_560.35.05_linux.run

wget -q --show-progress --progress=bar:force:noscroll -O cuda_installer.run $CUDA_URL
sudo sh cuda_installer.run --silent --toolkit --override

export CUDA_HOME=/usr/local/cuda-12.6

# Install PyTorch with CUDA 12.6
PYTORCH_INDEX_URL=https://download.pytorch.org/whl/cu126

pip3 install torch torchvision torchaudio --index-url $PYTORCH_INDEX_URL

# Additional packages for building extensions
pip3 install packaging ninja wheel setuptools setuptools-scm

Then install FlashAttention. For Hopper GPUs, install FlashAttention 3

git clone [email protected]:Dao-AILab/flash-attention.git
cd flash-attention/hopper
python setup.py install

For Ampere or earlier GPUs, install FlashAttention 2

pip3 install flash-attn

Install Python Dependencies 🐍

pip install -r requirements.txt

W&B Integration 📈

This project uses Weights & Biases for experiment tracking and metric visualization. Ensure you're logged in:

wandb login

Run Experiments

Quick Demo: Sudoku Solver 💻🗲

Train a master-level Sudoku AI capable of solving extremely difficult puzzles on a modern laptop GPU. 🧩

# Download and build Sudoku dataset
python dataset/build_sudoku_dataset.py --output-dir data/sudoku-extreme-1k-aug-1000  --subsample-size 1000 --num-aug 1000

# Start training (single GPU, smaller batch size)
OMP_NUM_THREADS=8 python pretrain.py data_path=data/sudoku-extreme-1k-aug-1000 epochs=20000 eval_interval=2000 global_batch_size=384 lr=7e-5 puzzle_emb_lr=7e-5 weight_decay=1.0 puzzle_emb_weight_decay=1.0

Runtime: ~10 hours on a RTX 4070 laptop GPU

Trained Checkpoints 🚧

To use the checkpoints, see Evaluation section below.

Full-scale Experiments 🔵

Experiments below assume an 8-GPU setup.

Dataset Preparation

# Initialize submodules
git submodule update --init --recursive

# ARC-1
python dataset/build_arc_dataset.py  # ARC offical + ConceptARC, 960 examples
# ARC-2
python dataset/build_arc_dataset.py --dataset-dirs dataset/raw-data/ARC-AGI-2/data --output-dir data/arc-2-aug-1000  # ARC-2 official, 1120 examples

# Sudoku-Extreme
python dataset/build_sudoku_dataset.py  # Full version
python dataset/build_sudoku_dataset.py --output-dir data/sudoku-extreme-1k-aug-1000  --subsample-size 1000 --num-aug 1000  # 1000 examples

# Maze
python dataset/build_maze_dataset.py  # 1000 examples

Dataset Visualization

Explore the puzzles visually:

Launch experiments

Small-sample (1K)

ARC-1:

OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py 

Runtime: ~24 hours

ARC-2:

OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py data_path=data/arc-2-aug-1000

Runtime: ~24 hours (checkpoint after 8 hours is often sufficient)

Sudoku Extreme (1k):

OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py data_path=data/sudoku-extreme-1k-aug-1000 epochs=20000 eval_interval=2000 lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0

Runtime: ~10 minutes

Maze 30x30 Hard (1k):

OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py data_path=data/maze-30x30-hard-1k epochs=20000 eval_interval=2000 lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0

Runtime: ~1 hour

Full Sudoku-Hard

OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py data_path=data/sudoku-hard-full epochs=100 eval_interval=10 lr_min_ratio=0.1 global_batch_size=2304 lr=3e-4 puzzle_emb_lr=3e-4 weight_decay=0.1 puzzle_emb_weight_decay=0.1 arch.loss.loss_type=softmax_cross_entropy arch.L_cycles=8 arch.halt_max_steps=8 arch.pos_encodings=learned

Runtime: ~2 hours

Evaluation

Evaluate your trained models:

OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 evaluate.py checkpoint=<CHECKPOINT_PATH>

Notes

Citation 📜

@misc{wang2025hierarchicalreasoningmodel,
      title={Hierarchical Reasoning Model}, 
      author={Guan Wang and Jin Li and Yuhao Sun and Xing Chen and Changling Liu and Yue Wu and Meng Lu and Sen Song and Yasin Abbasi Yadkori},
      year={2025},
      eprint={2506.21734},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2506.21734}, 
}

HRM (sapientinc/HRM)

Searching Database…

12.7K ·GitHub
这是什么?分层推理模型(HRM)的官方发布版本:一种受大脑启发的循环神经网络架构,仅用 2700 万参数就能达到很深的计算层次,只用约 1000 个样本训练,不需要预训练,也不需要思维链数据。
它是训练好的现成模型,还是需要我自己训练的代码?两者都有。这个仓库包含架构本身和完整的训练流程,作者还提供了 ARC-AGI-2、9x9 极难数独和 30x30 困难迷宫的已训练模型,你可以直接评估,不用自己训练。
我该怎么把它跑起来?先安装带 CUDA 的 PyTorch,因为仓库要编译 CUDA 扩展;然后安装 FlashAttention(Hopper 架构 GPU 用版本 3,Ampere 或更老的用版本 2),装好 Python 依赖,并登录 Weights & Biases 以便记录训练过程。
训练要多久,有什么要注意的坑?基于 1k 样本的数据集,在参考配置的 8 卡 GPU 环境下,数独 Extreme 大约 10 分钟训练完,30x30 迷宫约一小时,ARC 约 24 小时。小样本训练时准确率大约会有 ±2 个百分点的波动;在 Sudoku-Extreme 上要用早停,因为训练后期过拟合可能让训练和 Q-learning 变得不稳定。
我需要什么样的硬件才能训练?必须有 NVIDIA GPU——整套技术栈都依赖 PyTorch、CUDA 扩展和 FlashAttention。数独演示在一块 RTX 4070 笔记本 GPU 上约 10 小时就能训完,而完整规模的实验则假设使用 8 卡 GPU 环境。
它和大语言模型的推理方式有什么不同?它完全不使用思维链文本:HRM 在单次前向传播中,通过两个相互依赖的循环模块完成顺序推理——慢的高级模块负责抽象规划,快的低级模块负责细节计算——中间步骤没有任何显式监督。
它到底能解决哪些任务,效果怎么样?开箱即用有三类谜题式任务:9x9 极难数独、30x30 困难迷宫和抽象推理语料库(ARC)。在这些任务上,它的表现超过了上下文窗口长得多的大模型——数独和迷宫任务几乎做到满分。
我怎么看训练出的模型效果好不好?运行评估脚本,查看 Weights & Biases 里的 eval/exactaccuracy;ARC 任务有额外的步骤,最后在自带的 arceval.ipynb 笔记本里整理并查看结果。
我能自己看看训练用的谜题吗?可以——先生成数据集,再用浏览器打开 puzzlevisualizer.html,上传 data/ 里的文件夹,就能直观地浏览这些谜题。
Creator
sapientinc
Size
0.1 MB
Stars
12.7K
Updated
CommitMar 31
Age
1y
Commits · 90 Days