Please make this repo easier to use

[Tylersuard]

Not everyone has access to 4 H100 GPUs. Please either host a demo, put the model on huggingface, or create a Google Colab notebook so we can try it for ourselves. Thank you.

[alimaslax]

This is the cleanest ML code I seen in a while. You can run all her examples on a single RTX-6000. Look she even made wandb all setup for us. ya'll being too ungrateful. I killed the run after 4h cause I believed her. I wanna see what other grid like deterministic problems it can solve.

[zanngujjar]

Look at the pull request at #6 if you're having compatibility issues. This allowed me to train the model on 2 5060tis.

[gaoxin492]

I’ve successfully reproduced the results across four datasets.
The training curves and models are available at gaoxin492/TinyRecursiveModels
It turns out the results are quite close to those reported in the original paper:

Method	Params	Sudoku	Maze	ARC-1 (@2)	ARC-2 (@2)
TRM-Att	7M	77.71	78.70	41.00	3.33
TRM-MLP	5M	84.80	/	/	/

In the first training run, I didn’t plot the pass@k curves — I’ll include them in the next updated results. For the ARC experiments, I’ve increased the number of training epochs, and the performance is expected to improve further.

[thavidu]

Hi @gaoxin492 thanks for sharing your results! I'm seeing similar on my end- currently 76% of the way training and seeing the following pass@k for arc-agi-2
By the way, the huggingface cards you uploaded are currently empty I think- will you be uploading them later?
Also looking forward to seeing how your expanded epochs go- I'm training on 8xh100 unlike your 8xh200 so you're probably going to get there much faster than me :)

[gaoxin492]

@thavidu Yeah, for arc-agi-1 and arc-agi-2, I’ve increased the training steps since yesterday — hoping to reach the original reported results (44.6 and 7.8). Once the experiments finish, I’ll update the Hugging Face models and upload the full training curves as well.

Here’s a snapshot of the ongoing training curves — mine seem a bit smoother, probably because I’m using a larger global_batch_size:

[RonanKMcGovern]

Really nice work @gaoxin492 .

Those ARC results do seem meaningfully lower on II, and the pass at 2 does seem to be asymptoting.

How does your global batch size compare to the original? And how does your number of optimizer steps compare (with the run above, and with the run you are planning)?

Many thanks

[gaoxin492]

Thanks! @RonanKMcGovern

In the original setup, the global batch size was 768, while I set it to 4608 (6× larger) mainly to make full use of GPU memory.

For ARC-1, ARC-2, and Sudoku, I trained for 2× the number of epochs used by the authors — this allowed me to observe the full rise-and-fall trend of accuracy (i.e., up to overfitting). For Maze, I kept the original number of epochs, since the test accuracy was already clearly decreasing toward the end of training. You can find the corresponding training curves and details here: gaoxin492/TinyRecursiveModels

Because of the larger batch size, the gradient estimates are more accurate but typically yield slightly lower generalization performance, which likely explains why my reproduced results are a bit below the paper’s reported numbers.

[haruki-shirakami]

Thanks for the great work @gaoxin492.

Could you please clarify how you chose the final accuracy metrics to present? Did you choose the best performing checkpoint over the course of training, or the final checkpoint? Thanks.