HRM: Hierarchische Begründungsmodelle für komplexes Begründen

HRM (Hierarchical Reasoning Model) ist ein hierarchisches Denkmodell mit nur 27 Millionen Parametern, das für die Lösung komplexer Denkaufgaben im Bereich der künstlichen Intelligenz entwickelt wurde. Das Design des Modells ist inspiriert von der hierarchischen, mehrzeitskaligen Informationsverarbeitung des menschlichen Gehirns. Es führt sequentielle Denkaufgaben durch eine einzige Vorwärtspropagation ohne explizite Überwachung von Zwischenprozessen aus, und zwar über eine rekurrente neuronale Netzstruktur mit voneinander abhängigen High-Level-Modulen (zuständig für langsame, abstrakte Planung) und einem Low-Level-Modul (zuständig für schnelle, konkrete Berechnungen). HRM erreicht eine nahezu perfekte Leistung bei komplexen Aufgaben wie Sudoku und der Wegfindung in großen Labyrinthen mit nur 1000 Trainingsproben, ohne die Notwendigkeit eines Vortrainings oder von Chain-of-Thinking (CoT)-Daten, und übertrifft viele größere Modelle im Abstraction and Reasoning Corpus (ARC), einem wichtigen Benchmark für universelle KI. Der ARC ist der wichtigste Benchmark für allgemeine KI und übertrifft viele größere Modelle.

Funktionsliste

• Effizientes ReasoningEine neuartige rekurrente Architektur wird verwendet, um eine enorme Rechentiefe zu erreichen und gleichzeitig die Stabilität und Effizienz des Trainings zu erhalten.
• Keine Vorschulung erforderlichModelle können direkt aus einer kleinen Anzahl von Stichproben lernen, ohne dass ein umfangreicher Pre-Training-Prozess erforderlich ist.
• geringer DatenbedarfHohe Leistung bei komplexen Inferenzaufgaben kann mit nur 1000 Trainingsbeispielen erreicht werden.
• Zweimodulige StrukturEnthält ein High-Level-Modul für die abstrakte Planung und ein Low-Level-Modul für schnelle Berechnungen.
• Breite AnwendbarkeitLeistung bei einer Vielzahl komplexer Benchmarks, zum Beispiel:
  • Sudoku 9×9 Extrem
  • Labyrinth 30×30 Schwer
  • Abstraktions- und Reasoning-Korpus (ARC-AGI-2)
• die eigenen finanziellen Ressourcen zu erweiternDer Code steht auf GitHub zur Verfügung und bietet vortrainierte Modellprüfpunkte.

Hilfe verwenden

Für die Installation und Verwendung von HRM sind spezielle Hardware- und Softwareumgebungen erforderlich, vor allem NVIDIA-Grafikprozessoren mit CUDA-Unterstützung.Im Folgenden finden Sie eine detaillierte Beschreibung der Installation und Verwendung.

Vorbereitung der Umwelt

1. Installieren von CUDA:
   HRM stützt sich auf CUDA-Erweiterungen, und Sie müssen zunächst sicherstellen, dass der NVIDIA-Treiber und das CUDA-Toolkit auf Ihrem System installiert sind. Es wird empfohlen, CUDA Version 12.6 zu installieren.

   # 下载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:
   Installieren Sie das entsprechende PyTorch entsprechend der installierten CUDA-Version.

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

3. Installation des Build-Tools:
   Um die CUDA-Erweiterung kompilieren zu können, müssen einige zusätzliche Pakete installiert werden.

   pip3 install packaging ninja wheel setuptools setuptools-scm
   

4. FlashAttention installieren:
   Installieren Sie die entsprechende Version von FlashAttention für Ihr GPU-Modell.
   • Hopper-Architektur-GPUs (z. B. H100):

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

   • Ampere-Architektur oder frühere GPUs (z. B. RTX 30-Serie, 40-Serie):

     pip3 install flash-attn
     

5. Projektabhängigkeiten installieren:
   HRM-Repository klonen und installierenrequirements.txtAbhängigkeit in.

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

6. W&B-Integration (optional):
   Das Projekt verwendet Weights & Biases für die experimentelle Verfolgung. Wenn Sie den Trainingsprozess visualisieren möchten, besuchen Sie bitte W&B.

   wandb login
   

Schnellstart: Training einer Sudoku lösenden KI

Dies ist ein Beispiel für das Training eines schwierigen Sudoku-Lösungsmodells auf einer einzelnen Consumer-GPU (z. B. RTX 4070).

1. Aufbau des Datensatzes:
   Zunächst müssen Sie den Sudoku-Datensatz herunterladen und erstellen. Dieser Befehl generiert 1000 Beispiele von Sudoku auf hohem Niveau mit Datenerweiterung.

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

2. Beginn der Ausbildung:
   Verwenden Sie den folgenden Befehl, um das Training auf einer einzelnen GPU zu starten.

   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
   

   Bei der RTX 4070 dauert dieser Vorgang etwa 10 Stunden.

Groß angelegte Experimente

Für größere Experimente, wie z.B. ARC oder vollständige Sudoku-Datensätze, wird eine Multi-GPU-Umgebung (z.B. 8 Karten) empfohlen.

1. Initialisierung von Untermodulen:

   git submodule update --init --recursive
   

2. Vorbereiten des Datensatzes:
   Erstellen Sie je nach dem zu lösenden Problem den entsprechenden Datensatz.
   • ARC-2:

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

   • Labyrinth:

     python dataset/build_maze_dataset.py
     

3. Initiieren von Multi-GPU-Schulungen:
   ausnutzentorchrunInitiieren Sie verteiltes Training. Nehmen Sie das Beispiel des Trainings von Sudoku auf hohem Niveau (1000 Beispiele):

   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
   

   In einer 8-GPU-Umgebung dauert dieser Trainingsprozess etwa 10 Minuten.

Bewertungsmodell

Sie können vortrainierte Modell-Checkpoints verwenden oder Ihre eigenen trainierten Modelle bewerten.

1. Checkpoints herunterladen:
   Das offizielle Repository bietet vortrainierte Modelle für ARC-, Sudoku- und Labyrinth-Aufgaben.
2. Führen Sie das Auswertungsskript aus:
   Über die W&B-Schnittstelle betrachteteval/exact_accuracyMetriken. Für ARC-Aufgaben müssen zusätzliche Bewertungsskripte ausgeführt und die Ergebnisse mit Jupyter Notebook analysiert werden.

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

   Danach öffnen Sie diearc_eval.ipynbum die Ergebnisse abzuschließen und zu überprüfen.

caveat

• Die Genauigkeit des Lernens an kleinen Stichproben schwankt in der Regel um ±2 Punkte.
• Bei dem Tausend-Stichproben-Datensatz von Sudoku mit hohem Schwierigkeitsgrad kann es in der späten Phase des Trainings zu einer Überanpassung kommen, die zu instabilen Werten führt. Eine frühzeitige Abbruchstrategie wird empfohlen, wenn die Trainingsgenauigkeit nahe bei 100% liegt.

Anwendungsszenario

1. Kognitionswissenschaftliche Forschung
   Das Design von HRM ahmt den hierarchischen Informationsverarbeitungsmechanismus des menschlichen Gehirns nach und bietet ein berechenbares Modell für die Untersuchung menschlicher Planungs-, Denk- und Problemlösungsfähigkeiten und hilft dabei, den Weg zur Verwirklichung einer universellen KI zu erkunden.
2. Komplexe Planung und Terminierung
   In Bereichen wie Logistik, Roboterbahnplanung und automatisierte Produktionsplanung kann HRM schnell optimale oder nahezu optimale Lösungen finden, ohne auf große Datenmengen angewiesen zu sein, z. B. bei der Lösung großer Labyrinth-Pfadfindungsprobleme.
3. Spiel-KI
   Sie kann zur Entwicklung von KI eingesetzt werden, die in der Lage ist, komplexe Strategiespiele (z. B. Sudoku) zu lösen, wobei ihr effizientes Denkvermögen es ihr ermöglicht, die Spielregeln ohne vorheriges Training zu erlernen und zu beherrschen.
4. Lokalisierte KI-Anwendungen
   Aufgrund der geringen Anzahl von Modellparametern kann HRM auf lokalen Geräten mit begrenzten Ressourcen für Szenarien wie die Hausautomatisierung eingesetzt werden, um spezifische logische Schlussfolgerungen zu ziehen.

QA

1. Wie unterscheidet sich das HRM von der traditionellen Großsprachenmodellierung (LLM)?
   Anstatt sich auf die natürliche Sprache zu verlassen, führt HRM symbolische Manipulationen und Schlussfolgerungen durch zwei rekurrente Module mit unterschiedlichen Zeitskalen durch. Es erfordert kein umfangreiches Vortraining wie LLM und vermeidet die inhärente Fragilität, den hohen Datenbedarf und die hohe Latenz des Chain-of-Thinking-Ansatzes (CoT).
2. Welche Art von Hardware benötigen Sie, um HRM zu trainieren?
   Für das Training von HRM sind NVIDIA-GPUs mit CUDA-Unterstützung erforderlich. Es ist zwar möglich, auf einem einzelnen Consumer-GPU (z. B. einer RTX 3060 oder 4070) zu trainieren, aber das kann sehr lange dauern. Es wird offiziell empfohlen, für ein effizientes Training eine Multi-GPU-Umgebung (z. B. einen 8-Karten-Server) zu verwenden, z. B. dauert das Training eines Sudoku-Modells auf 8 GPUs nur 10 Minuten.
3. Ist die Menge der für das HRM erforderlichen Ausbildungsdaten hoch?
   Einer der größten Vorteile von HRM ist, dass es nur sehr wenige Trainingsdaten benötigt. Bei komplexen Aufgaben wie Sudoku, Labyrinthen und ARC erreicht es ein sehr hohes Leistungsniveau mit nur 1.000 Trainingsproben.
4. Wie lässt sich die in der Studie beschriebene Leistung reproduzieren?
   Aufgrund der stochastischen Natur des Lernens mit kleinen Stichproben können die Ergebnisse einer einzelnen Trainingssitzung leicht schwanken. Die Autoren der Studie stellen fest, dass die Standardabweichung der Genauigkeit bei der Sudoku-Aufgabe etwa 21 TP3 T beträgt. Um eine optimale Leistung zu erzielen, kann es notwendig sein, die Trainingsdauer fein abzustimmen und frühzeitig aufzuhören, bevor das Modell anfängt, überzupassen.