HRM：用于复杂推理的分层推理模型

HRM (Hierarchical Reasoning Model) 是一个仅有2700万参数的层级式推理模型，旨在解决人工智能领域中复杂的推理任务。该模型的设计灵感来源于人脑的层级式、多时间尺度的信息处理方式。 它通过一个高层模块（负责缓慢、抽象的规划）和一个低层模块（处理快速、具体的计算）相互依赖的循环神经网络结构，在没有明确中间过程监督的情况下，通过单次前向传播执行序列推理任务。 HRM无需预训练或思维链（CoT）数据，在仅使用1000个训练样本的情况下，即在数独和大型迷宫寻路等复杂任务上取得了接近完美的性能，并在通用人工智能关键基准“抽象与推理语料库”（ARC）上超越了许多更大规模的模型。

Funktionsliste

• Effizientes Reasoning：采用新颖的循环架构，在保持训练稳定性和效率的同时，获得了巨大的计算深度。
• 无需预训练：模型可以直接从少量样本中学习，无需大规模的预训练过程。
• 低数据需求：仅需1000个训练样本即可在复杂推理任务上达到高性能。
• 双模块结构：包含一个高层模块负责抽象规划，一个低层模块处理快速计算。
• 广泛适用性：在多种复杂的基准测试中表现出色，例如：
  • 高难度九宫格数独（Sudoku 9×9 Extreme）
  • 30×30迷宫路径寻找（Maze 30×30 Hard）
  • 抽象与推理语料库（ARC-AGI-2）
• die eigenen finanziellen Ressourcen zu erweitern：代码在GitHub上开源，并提供预训练好的模型检查点。

Hilfe verwenden

HRM的安装和使用需要特定的软硬件环境，主要是支持CUDA的NVIDIA GPU。以下是详细的安装和使用流程。

Vorbereitung der Umwelt

1. 安装CUDA:
   HRM依赖CUDA扩展，首先需要确保系统中已安装NVIDIA驱动和CUDA工具包。推荐安装CUDA 12.6版本。

   # 下载CUDA 12.6安装程序
   wget -q --show-progress --progress=bar:force:noscroll -O cuda_installer.run https://developer.download.nvidia.com/compute/cuda/12.6.3/local_installers/cuda_12.6.3_560.35.05_linux.run
   # 以静默模式安装
   sudo sh cuda_installer.run --silent --toolkit --override
   # 设置CUDA环境变量
   export CUDA_HOME=/usr/local/cuda-12.6
   

2. Installation von PyTorch:
   根据已安装的CUDA版本，安装对应的PyTorch。

   pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126
   

3. 安装构建工具:
   为了编译CUDA扩展，需要安装一些额外的包。

   pip3 install packaging ninja wheel setuptools setuptools-scm
   

4. 安装FlashAttention:
   根据你的GPU型号安装相应版本的FlashAttention。
   • Hopper架构GPU (例如 H100):

     git clone git@github.com:Dao-AILab/flash-attention.git
     cd flash-attention/hopper
     python setup.py install
     

   • Ampere架构或更早的GPU (例如 RTX 30系列, 40系列):

     pip3 install flash-attn
     

5. Projektabhängigkeiten installieren:
   克隆HRM仓库并安装requirements.txt中的依赖。

   git clone https://github.com/sapientinc/HRM.git
   cd HRM
   pip install -r requirements.txt
   

6. W&B集成（可选）:
   项目使用Weights & Biases进行实验跟踪。如果你需要可视化训练过程，请登录W&B。

   wandb login
   

快速上手：训练一个数独求解AI

这是一个在单张消费级GPU（例如RTX 4070）上训练高难度数独求解模型的示例。

1. 构建数据集:
   首先，需要下载并构建数独数据集。这个命令会生成1000个高难度数独样本，并进行数据增强。

   python dataset/build_sudoku_dataset.py --output-dir data/sudoku-extreme-1k-aug-1000 --subsample-size 1000 --num-aug 1000
   

2. Beginn der Ausbildung:
   使用以下命令在单个GPU上开始训练。

   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
   

   在RTX 4070上，这个过程大约需要10小时。

大规模实验

对于更大规模的实验，例如ARC或完整的数独数据集，推荐使用多GPU环境（例如8卡）。

1. Initialisierung von Untermodulen:

   git submodule update --init --recursive
   

2. Vorbereiten des Datensatzes:
   根据需要解决的问题，构建相应的数据集。
   • ARC-2:

     python dataset/build_arc_dataset.py --dataset-dirs dataset/raw-data/ARC-AGI-2/data --output-dir data/arc-2-aug-1000
     

   • 迷宫 (Maze):

     python dataset/build_maze_dataset.py
     

3. 启动多GPU训练:
   ausnutzentorchrun启动分布式训练。以训练高难度数独（1000样本）为例：

   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
   

   在8个GPU的环境下，这个训练过程大约只需要10分钟。

Bewertungsmodell

你可以使用预训练好的模型检查点，或者对自己训练的模型进行评估。

1. Checkpoints herunterladen:
   官方仓库提供了针对ARC、数独和迷宫任务的预训练模型。
2. Führen Sie das Auswertungsskript aus:
   通过W&B界面查看eval/exact_accuracy指标。对于ARC任务，需要额外运行评估脚本，并使用Jupyter Notebook进行结果分析。

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

   之后，打开arc_eval.ipynb来最终确定和检查结果。

caveat

• 小样本学习的准确率通常会有±2个点的浮动。
• 对于高难度数独的千样本数据集，训练后期可能会出现过拟合导致数值不稳定。建议在训练准确率接近100%时采用早期停止策略。

Anwendungsszenario

1. 认知科学研究
   HRM的设计模拟了人脑的层级化信息处理机制，为研究人类的规划、推理和解决问题能力提供了一个可计算的模型，有助于探索通用人工智能的实现路径。
2. 复杂规划与调度
   在物流、机器人路径规划、自动化生产调度等领域，HRM能够不依赖大量数据，快速找到最优或接近最优的解决方案，例如解决大型迷宫寻路问题。
3. 游戏AI
   可用于开发能够解决复杂策略游戏（如数独）的AI，其高效的推理能力使其能够在没有预训练的情况下学习并掌握游戏规则。
4. 本地化AI应用
   由于模型参数量小，HRM有潜力被部署在资源有限的本地设备上，用于家庭自动化等场景，执行特定的逻辑推理任务。

QA

1. HRM与传统大语言模型（LLM）有何不同？
   HRM不依赖自然语言进行推理，而是通过两个不同时间尺度的循环模块进行符号操作和推理。它不需要像LLM那样进行大规模的预训练，也避免了思维链（CoT）方法固有的脆弱性、高数据需求和高延迟等问题。
2. 训练HRM需要什么样的硬件？
   训练HRM需要支持CUDA的NVIDIA GPU。虽然可以在单张消费级GPU（如RTX 3060或4070）上进行训练，但这会花费较长时间。官方推荐使用多GPU环境（如8卡服务器）以实现高效训练，例如在8个GPU上训练数独模型仅需10分钟。
3. HRM的训练数据量要求高吗？
   不高。HRM的一大优势就是对训练数据的需求量极小。在数独、迷宫和ARC等复杂任务上，它仅使用1000个训练样本就达到了非常高的性能水平。
4. 如何复现论文中报告的性能？
   由于小样本学习的随机性，单次训练的结果可能会有轻微波动。论文作者指出，对于数独任务，准确率的标准差约为2%。为了达到最佳性能，可能需要微调训练时长，并在模型开始过拟合之前进行早期停止。