TRM-pretrain_arc2_rearc_100k

Tiny Recursive Model Trained on Concept ARC ARC AGI II Training + Re-ARC

Trelis is open to compute sponsorship, specifically for 4xH100 SXM or 8xH100 SXM. Sponsors will be credited on related posts and Trelis Youtube videos. Past sponsors include Runpod (affiliate link incl. sign-up bonus) and Lambda Labs. E-mail arc at trelis dot com

For machine learning tutorials and tools, see the Trelis Youtube Channel, see also Trelis on X

TL;DR: After 200k epochs of training on ARC AGI II training, Concept ARC and Re-arc datasets, the model achieves ~15% on ARC AGI II Public Evaluation tasks, meaningfully above previously reported scores of ~8% after training for 100k epochs on ARC AGI II training and Concept ARC.

Models:

• The 100k epochs model is here
• The 200k epochs model is here (ctd pretraining of the 100k model, using the same dataset build, i.e. augmentations) Note that, unlike previous model repos, these include the built datasets, so one does not need to re-initialise embeddings to do continued pre-training on.

CREDIT: All of this work builds on the HRM and TRM repos and papers.

Training Notes

The Github Repo is fully open source here. Data is all in weights and biases here.

Dataset Preparation

This training run includes:

• Michael Hodel's Re-arc dataset of 400 ARC AGI I training tasks, including 1000 examples per task, with one task removed because it contains at least one grid larger than 30x30. 399 total tasks. Note that there are 1,000 train examples per task and no test examples.
• ARC AGI II Training Data of 1,000 tasks, deduped by ARC AGI I training tasks to leave 600 tasks. Both train and test examples are used. There are an average of about 3 train examples per task and 1 test example per task.
• ARC AGI II Evaluation Data of 120 tasks. Note that only the train examples are included in pre-training.

Evaluation is conducted only on the test examples of the 120 ARC AGI II public evaluation data.

Other Notes

• HRM/TRM fills batches on a per task basis (after selecting a task variant). With a global batch size of 768 and 1,000 examples per Re-arc task (compared to no more than 10+3 for ARC AGI II training or Concept ARC examples), that would result in batches full of Re-arc tasks that cause periodicity in the loss curve during training. To mitigate this, a cap of 5 examples per task was applied (sampled at random) per batch during training.
• Re-arc task examples come from a generator function for original ARC AGI I training tasks. This generator function varies inputs at random to create augmented examples. By contrast the HRM/TRM code applies simpler, but systematic augmentations, when building the dataset. As such, Re-arc tasks contain augmented examples (at random), which are then subjected to HRM/TRM augmentation. By contrast, ARC AGI II and Concept ARC tasks each contain tasks that are more self consistent, and are then augmented. Applying 1,000x augmentations to the re-arc data would have been excessive, and so only 3 augmentations were used for re-arc tasks, to give roughly 400x1000x3 =1,200,000 examples, compared to roughly 750x3x1000 = 2,250,000 examples across ARC AGI II training and Concept ARC. There is a further nuance here in that each ARC variant gets its own embedding, that should be unique to the task but also to the systematic variant (rotation/flip, re-colour). Re-arc examples mix such variants (and more) and so the task embedding for re-arc tasks must aggregate more information on variants. It is hard to know whether this provides beneficial regularisation or over-smooths.
• Beware, when building the dataset - since there are two bespoke splits (owing to the use of a different number of augmentations for each) that one must pass puzzle_identifiers_start.

Architecture

The same was used as the original TRM work, except for 4 lower cycles instead of 6 stated in the paper.

Results

Pre-training Results

Pass@2 (competition scoring) and Pass@1000 evaluation scores versus optimizer steps (not epochs) are shown below, the results of 100k epochs of pre-training on ARC AGI II training and Concept ARC alone are included.

Commentary:

• Pass@2 does not clearly asymptote at 200k epochs, suggesting further gains.
• Pass@1000k continues to increase at 200k epochs, although there are signs of asymptoting. This suggests further gains from further pre-training.
• Allowing for noise, it is not clear that adding the re-arc dataset improves performance over pre-training solely on ARC AGI II training and Concept ARC.

Fig. 1: Pass@2 Evaluation Score vs Optimizer Steps

Fig. 2: Pass@1000 Evaluation Score vs Optimizer Steps

Post-training Results

Although late submissions are still being accepted (although not eligible for the competition) for the ARC Prize 2025, there is no longer access to L4 GPUs, meaning it is not possible to save a new run. This means it is not possible to measure post-training performance on the ARC AGI II semi-private dataset.

At 100k epochs of training, it was possible to get 6.67% semi-private score with a ~10% public score pre-trained model. Very speculatively, perhaps ~10% could be achieved on semi-private with a pre-trained model scoring 15% on the public evaluation set.

Future Work

I have tried a large number of hyperparameter changes to the model without a clear improvement (although one must run for a long time to see such improvement).

Obvious next steps are:

• Training for longer. Potentially 20%+ is within sight in 100k epochs more.
• Training a larger model. Potentially a larger model is stronger on an iso-compute basis. It becomes more difficult to inference in Kaggle, but there is still a window there for headroom.