Views
No views yet

1pip install --upgrade pip wheel setuptools
2pip install --pre --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu126 # install torch based on your cuda version
3pip install -r requirements.txt # install requirements
4pip install --no-cache-dir --no-build-isolation adam-atan2
5wandb login YOUR-LOGIN # login if you want the logger to sync results to your Weights & Biases (https://wandb.ai/)1# ARC-AGI-1
2python -m dataset.build_arc_dataset \
3 --input-file-prefix kaggle/combined/arc-agi \
4 --output-dir data/arc1concept-aug-1000 \
5 --subsets training evaluation concept \
6 --test-set-name evaluation
7
8# ARC-AGI-2
9python -m dataset.build_arc_dataset \
10 --input-file-prefix kaggle/combined/arc-agi \
11 --output-dir data/arc2concept-aug-1000 \
12 --subsets training2 evaluation2 concept \
13 --test-set-name evaluation2
14
15## Note: You cannot train on both ARC-AGI-1 and ARC-AGI-2 and evaluate them both because ARC-AGI-2 training data contains some ARC-AGI-1 eval data
16
17# Sudoku-Extreme
18python dataset/build_sudoku_dataset.py --output-dir data/sudoku-extreme-1k-aug-1000 --subsample-size 1000 --num-aug 1000 # 1000 examples, 1000 augments
19
20# Maze-Hard
21python dataset/build_maze_dataset.py # 1000 examples, 8 augments1run_name="pretrain_att_arc1concept_4"
2torchrun --nproc-per-node 4 --rdzv_backend=c10d --rdzv_endpoint=localhost:0 --nnodes=1 pretrain.py \
3arch=trm \
4data_paths="[data/arc1concept-aug-1000]" \
5arch.L_layers=2 \
6arch.H_cycles=3 arch.L_cycles=4 \
7+run_name=${run_name} ema=True
81run_name="pretrain_att_arc2concept_4"
2torchrun --nproc-per-node 4 --rdzv_backend=c10d --rdzv_endpoint=localhost:0 --nnodes=1 pretrain.py \
3arch=trm \
4data_paths="[data/arc2concept-aug-1000]" \
5arch.L_layers=2 \
6arch.H_cycles=3 arch.L_cycles=4 \
7+run_name=${run_name} ema=True
81run_name="pretrain_mlp_t_sudoku"
2python pretrain.py \
3arch=trm \
4data_paths="[data/sudoku-extreme-1k-aug-1000]" \
5evaluators="[]" \
6epochs=50000 eval_interval=5000 \
7lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0 \
8arch.mlp_t=True arch.pos_encodings=none \
9arch.L_layers=2 \
10arch.H_cycles=3 arch.L_cycles=6 \
11+run_name=${run_name} ema=True
12
13run_name="pretrain_att_sudoku"
14python pretrain.py \
15arch=trm \
16data_paths="[data/sudoku-extreme-1k-aug-1000]" \
17evaluators="[]" \
18epochs=50000 eval_interval=5000 \
19lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0 \
20arch.L_layers=2 \
21arch.H_cycles=3 arch.L_cycles=6 \
22+run_name=${run_name} ema=True1run_name="pretrain_att_maze30x30"
2torchrun --nproc-per-node 4 --rdzv_backend=c10d --rdzv_endpoint=localhost:0 --nnodes=1 pretrain.py \
3arch=trm \
4data_paths="[data/maze-30x30-hard-1k]" \
5evaluators="[]" \
6epochs=50000 eval_interval=5000 \
7lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0 \
8arch.L_layers=2 \
9arch.H_cycles=3 arch.L_cycles=4 \
10+run_name=${run_name} ema=True1@misc{jolicoeurmartineau2025morerecursivereasoningtiny,
2 title={Less is More: Recursive Reasoning with Tiny Networks},
3 author={Alexia Jolicoeur-Martineau},
4 year={2025},
5 eprint={2510.04871},
6 archivePrefix={arXiv},
7 primaryClass={cs.LG},
8 url={https://arxiv.org/abs/2510.04871},
9}1@misc{wang2025hierarchicalreasoningmodel,
2 title={Hierarchical Reasoning Model},
3 author={Guan Wang and Jin Li and Yuhao Sun and Xing Chen and Changling Liu and Yue Wu and Meng Lu and Sen Song and Yasin Abbasi Yadkori},
4 year={2025},
5 eprint={2506.21734},
6 archivePrefix={arXiv},
7 primaryClass={cs.AI},
8 url={https://arxiv.org/abs/2506.21734},
9}