HRM：複雑な推論のための階層的推論モデル

HRM（Hierarchical Reasoning Model）は、人工知能分野における複雑な推論タスクを解くために設計された、わずか2700万個のパラメータを持つ階層的推論モデルである。このモデルの設計は、人間の脳の階層的でマルチタイムスケールな情報処理にヒントを得ている。 HRMは、相互依存的な高レベルモジュール（低速で抽象的なプランニングを担当）と低レベルモジュール（高速で具体的な計算を担当）を持つリカレントニューラルネットワーク構造を介して、明示的な中間プロセスの監視なしに、単一の順伝播により逐次推論タスクを実行する。 HRMは、事前学習や思考連鎖（CoT）データを必要とせず、わずか1000個の学習サンプルで、数独や大規模迷路の経路探索などの複雑なタスクにおいて完璧に近い性能を達成し、汎用AIの重要なベンチマークである抽象化推論コーパス（ARC）において、多くの大規模モデルを凌駕した。ARCは汎用AIの重要なベンチマークであり、多くの大規模モデルを凌駕する。

機能一覧

• 効率的な推論このように、リカレントアーキテクチャを用いることで、学習の安定性と効率性を維持しながら、圧倒的な計算量を得ることができる。
• 事前トレーニング不要モデルは、大規模な事前学習プロセスなしに、少数のサンプルから直接学習することができます。
• 低データ需要複雑な推論タスクにおいて、わずか1000個のトレーニングサンプルで高いパフォーマンスを達成することができる。
• 2モジュール構造抽象的なプランニングを行う高水準モジュールと、高速計算を行う低水準モジュールを含む。
• 幅広い適用性例えば、様々な複雑なベンチマークで好成績を収めている：
  • ナンプレ9×9エクストリーム
  • 迷路 30×30 ハード
  • 抽象化と推論コーパス（ARC-AGI-2）
• ざいげんをひろげるコードはGitHubでオープンソース化されており、事前にトレーニングされたモデルのチェックポイントを提供しています。

ヘルプの使用

HRMのインストールと使用には、特定のハードウェアとソフトウェア環境、主にCUDAをサポートするNVIDIA GPUが必要です。

環境準備

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. 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をインストールしてください。
   • ホッパー・アーキテクチャ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. プロジェクトの依存関係をインストールする:
   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つで難しい数独の解答モデルをトレーニングした例です。

1. データセットの構築:
   まず、数独データセットをダウンロードしてビルドする必要がある。このコマンドは、データ補強を行った高レベルの数独の1000サンプルを生成する。

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

2. トレーニング開始:
   単一の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. サブモジュールの初期化:

   git submodule update --init --recursive
   

2. データセットの準備:
   解決すべき問題に応じて、適切なデータセットを構築する。
   • ARC-2:

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

   • 迷路:

     python dataset/build_maze_dataset.py
     

3. マルチGPUトレーニングの開始:
   利用するtorchrun分散トレーニングを開始する。高レベルの数独（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
   

   GPU8基の環境では、このトレーニングにかかる時間は約10分。

アセスメントモデル

事前に訓練されたモデルのチェックポイントを使用することも、自分で訓練したモデルを評価することもできます。

1. ダウンロード Checkpoints:
   公式リポジトリでは、ARC、数独、迷路タスクの訓練済みモデルを提供しています。
2. 評価スクリプトの実行:
   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を最終確認し、結果をチェックする。

ほら

• 小サンプル学習の精度は通常±2ポイント変動する。
• 高難易度数独の千サンプルデータセットでは、学習後期にオーバーフィッティングが起こり、値が不安定になる可能性がある。学習精度が100%に近い場合は、早期の停止を推奨する。

アプリケーションシナリオ

1. 認知科学研究
   HRMの設計は、人間の脳の階層的情報処理メカニズムを模倣しており、人間の計画、推論、問題解決能力を研究するための計算可能なモデルを提供し、汎用AIの実現への道を探る一助となる。
2. 複雑なプランニングとスケジューリング
   ロジスティクス、ロボット経路計画、自動生産スケジューリングなどの分野において、HRMは、大規模な迷路経路探索問題を解くように、大量のデータに依存することなく、最適解または最適解に近い解を迅速に見つけることができる。
3. ゲームAI
   これは、複雑な戦略ゲーム（例：数独）を解くことができるAIの開発に使用することができ、その効率的な推論により、事前訓練なしでゲームのルールを学習し、習得することができる。
4. ローカライズされたAIアプリケーション
   モデルパラメータの数が少ないため、HRMは、特定の論理的推論タスクを実行するホームオートメーションなどのシナリオにおいて、リソースが限られたローカルデバイスに展開できる可能性がある。

品質保証

1. HRMは従来のラージ・ランゲージ・モデリング（LLM）とどう違うのか？
   HRMは、推論を自然言語に依存する代わりに、異なる時間スケールを持つ2つのリカレントモジュールを通じて、記号操作と推論を行う。LLMのような大規模な事前学習を必要とせず、Chain of Thought (CoT)アプローチ特有の脆弱性、高データ要件、高レイテンシを回避することができる。
2. HRMのトレーニングにはどのようなハードウェアが必要ですか？
   HRMのトレーニングには、CUDAをサポートするNVIDIA GPUが必要です。単一のコンシューマーGPU（RTX 3060や4070など）でトレーニングすることは可能ですが、これには長い時間がかかります。効率的なトレーニングのためには、マルチGPU環境（例：8カードサーバー）を使用することが公式に推奨されています（例：数独モデルを8GPUでトレーニングする場合、わずか10分しかかかりません）。
3. HRMに必要なトレーニングデータの量は多いのか？
   高くない。HRMの大きな利点の1つは、学習データをほとんど必要としないことである。数独、迷路、ARCのような複雑なタスクでは、わずか1,000のトレーニングサンプルで非常に高いレベルのパフォーマンスを達成する。
4. 論文で報告されたパフォーマンスを再現するには？
   小サンプル学習の確率的性質のため、1回のトレーニングセッションの結果はわずかに変動する可能性がある。この論文の著者は、数独タスクの場合、精度の標準偏差は約21 TP3 Tであると指摘している。最適なパフォーマンスを得るためには、訓練時間を微調整し、モデルがオーバーフィットし始める前に早期に停止する必要があるかもしれない。