HRM: modelos de raciocínio hierárquico para raciocínio complexo

O HRM (Hierarchical Reasoning Model) é um modelo de raciocínio hierárquico com apenas 27 milhões de parâmetros projetado para resolver tarefas de raciocínio complexas no campo da inteligência artificial. O design do modelo é inspirado no processamento de informações hierárquico e em várias escalas de tempo do cérebro humano. Ele executa tarefas de raciocínio sequenciais por meio de uma única propagação direta sem supervisão explícita de processos intermediários, por meio de uma estrutura de rede neural recorrente com módulos interdependentes de alto nível (responsáveis pelo planejamento lento e abstrato) e um módulo de baixo nível (que lida com cálculos rápidos e concretos). O HRM alcança um desempenho quase perfeito em tarefas complexas, como Sudoku e localização de caminhos em grandes labirintos, com apenas 1.000 amostras de treinamento, sem a necessidade de pré-treinamento ou dados de cadeia de raciocínio (CoT), e supera muitos modelos maiores no Abstraction and Reasoning Corpus (ARC), uma referência importante para IA de uso geral. O ARC é a principal referência para a IA de uso geral, superando muitos modelos de grande escala.

Lista de funções

• Raciocínio eficienteArquitetura recorrente: Uma nova arquitetura recorrente é usada para obter uma enorme profundidade computacional, mantendo a estabilidade e a eficiência do treinamento.
• Não é necessário treinamento prévioModelos: Os modelos podem aprender diretamente com um pequeno número de amostras sem um processo de pré-treinamento em grande escala.
• baixa demanda de dadosO que significa que o alto desempenho em tarefas de inferência complexas pode ser obtido com apenas 1.000 amostras de treinamento.
• Estrutura de dois módulosMódulo de alto nível: contém um módulo de alto nível para planejamento abstrato e um módulo de baixo nível para lidar com a computação rápida.
• Ampla aplicabilidadeDesempenho em uma variedade de benchmarks complexos, por exemplo:
  • Sudoku 9×9 Extreme
  • Labirinto 30×30 Difícil
  • Corpus de abstração e raciocínio (ARC-AGI-2)
• expandir seus recursos financeirosO código é de código aberto no GitHub e fornece pontos de verificação de modelos pré-treinados.

Usando a Ajuda

A instalação e o uso do HRM requerem ambientes específicos de hardware e software, principalmente GPUs NVIDIA com suporte a CUDA.

Preparação ambiental

1. Instalação do CUDA:
   O HRM depende de extensões CUDA e, primeiro, você precisa se certificar de que o driver NVIDIA e o kit de ferramentas CUDA estejam instalados em seu sistema. Recomenda-se a instalação da versão 12.6 do CUDA.

   # 下载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. Instalando o PyTorch:
   Instale o PyTorch correspondente de acordo com a versão CUDA instalada.

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

3. Instalação da ferramenta de compilação:
   Para compilar a extensão CUDA, alguns pacotes adicionais precisam ser instalados.

   pip3 install packaging ninja wheel setuptools setuptools-scm
   

4. Instalar o FlashAttention:
   Instale a versão apropriada do FlashAttention para seu modelo de GPU.
   • GPUs com arquitetura Hopper (por exemplo, H100):

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

   • Arquitetura Ampere ou GPUs anteriores (por exemplo, RTX série 30, série 40):

     pip3 install flash-attn
     

5. Instalar as dependências do projeto:
   Clonar o repositório HRM e instalarrequirements.txtDependência em.

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

6. Integração de W&B (opcional):
   O projeto usa o Weights & Biases para rastreamento experimental. Se você precisar visualizar o processo de treinamento, visite W&B.

   wandb login
   

Início rápido: treinamento de uma IA para resolução de Sudoku

Este é um exemplo de treinamento de um modelo difícil de resolução de Sudoku em uma única GPU de consumidor (por exemplo, RTX 4070).

1. Construção do conjunto de dados:
   Primeiramente, você precisa fazer o download e criar o conjunto de dados do Sudoku. Esse comando gera 1.000 amostras de Sudoku de alto nível com aumento de dados.

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

2. Iniciar o treinamento:
   Use o seguinte comando para iniciar o treinamento em uma única 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
   

   Na RTX 4070, esse processo leva cerca de 10 horas.

Experimentos em grande escala

Para experimentos maiores, como ARC ou conjuntos de dados completos de Sudoku, recomenda-se um ambiente com várias GPUs (por exemplo, 8 placas).

1. Inicialização de submódulos:

   git submodule update --init --recursive
   

2. Preparação do conjunto de dados:
   Dependendo do problema a ser resolvido, crie o conjunto de dados apropriado.
   • ARC-2:

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

   • Labirinto:

     python dataset/build_maze_dataset.py
     

3. Iniciar o treinamento em várias GPUs:
   fazer uso detorchrunIniciar o treinamento distribuído. Veja o exemplo do treinamento de Sudoku de alto nível (1.000 amostras):

   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
   

   Em um ambiente de 8 GPUs, esse processo de treinamento leva cerca de 10 minutos.

modelo de avaliação

Você pode usar pontos de verificação de modelos pré-treinados ou avaliar seus próprios modelos treinados.

1. Baixar pontos de controle:
   O repositório oficial fornece modelos pré-treinados para as tarefas ARC, Sudoku e Maze.
2. Executar o script de avaliação:
   Visualizado por meio da interface W&Beval/exact_accuracyMétricas. Para tarefas ARC, scripts de avaliação adicionais precisam ser executados e os resultados analisados usando o Jupyter Notebook.

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

   Depois disso, abra oarc_eval.ipynbpara finalizar e verificar os resultados.

advertência

• A precisão do aprendizado de amostras pequenas normalmente flutua em ±2 pontos.
• No caso do conjunto de dados de mil amostras de Sudoku de alta dificuldade, pode ocorrer um ajuste excessivo no estágio final do treinamento, resultando em valores instáveis. Recomenda-se uma estratégia de parada antecipada quando a precisão do treinamento estiver próxima de 100%.

cenário do aplicativo

1. Pesquisa em Ciências Cognitivas
   O design do HRM imita o mecanismo hierárquico de processamento de informações do cérebro humano, fornecendo um modelo computável para estudar o planejamento humano, o raciocínio e as habilidades de resolução de problemas, além de ajudar a explorar o caminho para a realização da IA de propósito geral.
2. Planejamento e programação complexos
   Em áreas como logística, planejamento de caminhos de robôs e programação de produção automatizada, o HRM é capaz de encontrar rapidamente soluções ideais ou quase ideais sem depender de grandes quantidades de dados, como a solução de grandes problemas de localização de caminhos em labirintos.
3. IA do jogo
   Ele pode ser usado para desenvolver IAs capazes de resolver jogos de estratégia complexos (por exemplo, Sudoku), em que seu raciocínio eficiente permite que ele aprenda e domine as regras do jogo sem treinamento prévio.
4. Aplicativos de IA localizados
   Devido ao pequeno número de parâmetros do modelo, o HRM tem o potencial de ser implantado em dispositivos locais com recursos limitados para cenários como a automação residencial para executar tarefas específicas de raciocínio lógico.

QA

1. Como o HRM difere do LLM (Large Language Modelling) tradicional?
   Em vez de depender da linguagem natural para o raciocínio, o HRM realiza a manipulação simbólica e o raciocínio por meio de dois módulos recorrentes com escalas de tempo diferentes. Ele não exige pré-treinamento em larga escala, como o LLM, e evita a fragilidade inerente, os altos requisitos de dados e a alta latência da abordagem de cadeia de raciocínio (CoT).
2. Que tipo de hardware é necessário para treinar o HRM?
   O treinamento de HRM requer GPUs NVIDIA com suporte a CUDA e, embora seja possível treinar em uma única GPU de consumo (como uma RTX 3060 ou 4070), isso pode levar muito tempo. É oficialmente recomendado o uso de um ambiente com várias GPUs (por exemplo, servidor de 8 placas) para um treinamento eficiente, por exemplo, o treinamento de um modelo de Sudoku em 8 GPUs leva apenas 10 minutos.
3. A quantidade de dados de treinamento necessários para o HRM é alta?
   Uma das principais vantagens do HRM é que ele requer pouquíssimos dados de treinamento. Em tarefas complexas, como Sudoku, labirintos e ARC, ele atinge um nível muito alto de desempenho usando apenas 1.000 amostras de treinamento.
4. Como reproduzir o desempenho relatado no artigo?
   Devido à natureza estocástica do aprendizado de pequenas amostras, os resultados de uma única sessão de treinamento podem variar ligeiramente. Os autores do artigo observam que, para a tarefa de Sudoku, o desvio padrão da precisão é de cerca de 21 TP3 T. Para obter o desempenho ideal, pode ser necessário ajustar a duração do treinamento e interrompê-lo antes que o modelo comece a se ajustar demais.