Training Guide: HRM SRAM/DRAM Learning Performance

Use this guide to train the memory-tiered HRM on real datasets and verify its reasoning accuracy.

1. Dataset Generation

The model needs to learn how to solve puzzles. First, generate a training/test set (e.g., Sudoku).

# Activate environment
source venv/bin/activate

# Build a small Sudoku dataset (1000 examples)
python dataset/build_sudoku_dataset.py \
    --output-dir data/sudoku-1k \
    --subsample-size 1000 \
    --num-aug 10

2. Start Training

Point the trainer to the tiered architecture configuration.

Train the Tiered Model (New)

python pretrain.py \
    arch=hrm_tiered \
    data_path=data/sudoku-1k \
    epochs=1000 \
    global_batch_size=384

Train the Baseline Model (Comparison)

python pretrain.py \
    arch=hrm_v1 \
    data_path=data/sudoku-1k \
    epochs=1000 \
    global_batch_size=384

3. Verify Reasoning Accuracy

Once training is complete (or during training), check the accuracy metrics in your W&B dashboard or via the evaluation script:

# Replace with the path to your generated checkpoint
python evaluate.py checkpoint=checkpoints/Sudoku_ACT-torch/YOUR_RUN/step_1000

What to look for:

eval/exact_accuracy: Fraction of puzzles solved perfectly.
eval/steps: Average number of reasoning steps taken by the ACT (Adaptive Computation Time) module.
eval/q_halt_accuracy: How well the model learns when to stop thinking.

4. Training on ARC (Artificial General Intelligence Benchmark)

To train on the more complex ARC-AGI-2 dataset:

python dataset/build_arc_dataset.py --output-dir data/arc-2
python pretrain.py arch=hrm_tiered data_path=data/arc-2