44% on ARC-AGI-1 in 67 cents

I trained a small transformer from scratch in 1.5hrs on a 5090 Beats many LLMs, and scores the same as TRM/HRMThis is an upgrade to my previous model Faster, better, cheaper and still open source.Also gets 7% on ARC-2Discussion on Twitter, Code on github Performance on ARC-1 public eval. I only compare against models that do similar test time training This is the 3rd blog in a series of works on ARC-AGI. Prev: Blog 2, Blog 1.Many ppl thought the prev result was impossible. It got attention from top researchers and went viral on X. Eg: Discussions by Lucas Beyer, Jeremy Howard, Rohan Anil, and comments by many others.Why work on this?I think sample efficiency is the most important problem in AI today and I want to solve it.The intention behind this work is to (1) find the limits of sample efficiency when restricted to transformers / today’s deep learning methods and (2) reduce costs so iteration is much faster and cheaper.ARC is a great benchmark to test this:Very few samples (only a 1000 puzzles) in a high dimensional spaceIts a metalearning benchmark, so each puzzle uses a different rule, with some common conceptsVery few priors needed: every concept needed in the eval set is present in the train setIt is incredibly easy for humans to solve, and accessible to even poor AI researchersBenchmark is still unsaturated (for data efficiency, ignore LLMs and approaches that use tons of synthetic data or human inductive biases)Next, I’ll work on new research ideas to break these limits. I’ll try to keep costs low so that anyone in the world can work on this.Tech detailsHow does it work?The overall approach is similar to last time (full technical details here), but I added a bunch of upgrades. Here’s a quick summary of the approach:Each input-output pair is converted to a sequence of tokens. These sequences are autoregressively trained on by a small transformer. This is done from scratch at test time on both the train set and eval set puzzles (test labels hidden).To enable cross-task learning, each puzzle is given a separate additive embedding (learnt). Since each sequence has two 2D grids, positional are learnt using 3D RoPE embeddings.The sequences are augmented with color and dihedral permutations. During inference, the test inputs are augmented, and the inverse aug is applied on the outputs produced. The 2 most common outputs are submitted (AAIVR).Changes since last timeThe main goal was to find improvements to the architecture / algorithm that improve the sample efficiency of the model.The biggest increases in scores were due toModern architecture (SwiGlu instead of GELU, RMSnorm not layernorm, etc.)More data diversity, better shuffling of datascaling up: 8 layers instead of 4,Biggest decreases in cost were due to:Way fewer augmentations (more sample efficient!)AdamW -> Normuonflash attention with varlen training + flex attention kernels for inferenceA major change is that I don’t train on input tokens anymore. This means the loss function only includes output tokens (which makes the approach supervised). This. performs slightly better 40% $\to$ 44% but I don’t understand why. Perhaps finite model capacityI also increased the training data by adding the non-overlapping tasks from ARC-2. I did this very carefully to ensure no leakage. You can remove the extra data if you don’t like it and it will still score ~40%, but it will need ~double the compute.Context: ARC-2 contains 773 ARC-1 puzzles and 347 new puzzles. Most eval puzzles of ARC-1 are repeated, so if you naively train on ARC-2, then its a dataleak and you will score 100%. I avoid this by carefully filtering out the 773 repeated puzzles (so no leak!)There are many other changes that gave incremental improvements in performance or speed. Find the full list of changes here.Interesting behaviourSince I am no longer training on inputs, this approach is now supervised. What’s weird is that the test loss is now worse, yet it scores better! Also it is more stable and there’s less variance in scores. Many ppl today are working on sample efficiency by aiming for the lowest val loss on a small dataset. I think that’s great, but this points out a failure mode in such an approachI do think the unsupervised style training will be better in some scenarios, and I am evaluating this.Before NorMuon, I tried vanilla Muon. Obviously it trained much faster than AdamW, but the loss (and scores) would loiter at the end instead of converging. I found that cranking down the momentum and/or LR drastically at this point helped, but I didn’t want to make manually changes like this. When I switched to NorMuon, the problem disappearedAblationsThe biggest contribution to performance seems to be good representations (3D RoPE + per-task embedding). Removing 3D RoPE or the per-task embedding gives a steep drop. Both ablations saturate at 25%Training on inputs performs slightly worse -> ~39%Restricting training set to ARC-1+ConceptARC only performs about the same: ~40%Switching from 3D RoPE to 1D drops score to ~24%Removing the per-task embeddings drops score to ~24%Running the model CompressARC style (training from scratch on each task separately, and unsupervised), gives a drops performance down to ~18%CompressARC but supervised gets ~15% Finding the best scores on other ablations. Comparing costs makes little sense here as all but the first ablation requires a lot more compute How can others contribute?The code is open source. Feel free to modify it and improve score or reduce cost. (Pls don’t increase training data)Try reaching 65% – you won’t need many modifications. Evidence: I took the union of all solved tasks from multiple runs, and got 55%. Also a bunch of other tasks are “almost” solved. Some ideas:RoPE mixes positional and content information, which probably worsens performance. PoPE should perform on par or better. Or maybe invent a new pos embeddingThe architecture can definitely be modernised furtherCosts can probably be reduced 10x with handmade GPU code. There are architectural changes that can also do this.Lastly, figure out how to remove data augmentations. (I hate that I used it, ignore everyone who thinks its okay). There are a few obvious ways to do so, but the challenge is keeping training costs low.MiscTBH, I didn’t expect to reach 45% with just the transformer, I thought this would need new ideas. I certainly didn’t expect to reach it at such low costs/flops. The ablations show that a surprising amount of perfomance is retained even without augmentations or synthetic data. Now I’m pretty sure 65% can be reached within the transformer frameworkI don’t understand why others didn’t figure this out. Its just a transformer with the most obvious representation. This benchmark has been open for 6 years, was high profile, and had a million dollar prize! Maybe researchers underestimate deep learning? Maybe the cost of experimentation was high enough that they couldn’t run ablations properly? Blindsided by LLMs or using harnesses?AppendixPrev criticism/validation on my approach from famous researchersMy old result went viral on X and many experienced researchers debated about it, both for and against. Threads by Jeremy, Lucas, Susan, Andreas, Yoav, and many more. I’m listing all the criticisms here with my answers.Training on the eval puzzles is cheating / “training on test”No this is false. “Training on test” specifically means training on the labels of test data. The labels were not trained on.Also, ARC is a metalearning benchmark, so you’re supposed to learn from the eval puzzles.Jargon: ARC has a set of train puzzles and a set of eval puzzles. Each puzzle has example pairs and test pairs. A pair consists of an input grid + output grid.The ARC, the label is only the test pair’s output grid in an eval puzzle.These labels were not trained on. They are hidden. You can delete it beforehand if you wish Training on the inputs of eval puzzles leaks informationNo, this is false. Such an approach is called transductive reasoning and has been studied since the time of Vapnik.Also, this dogma of ignoring eval inputs doesn’t make sense in a world trying to solve continual learningOther approaches train a metalearning algorithm and then deploy it to learn by running a CoT or by modifying latents through a recurrent loo. My approach or what I did here is directly metalearn by modifying the weights of a single forward function is no different than learning byNote: in the new 44% result, training on inputs has been removed as it scores slightly worseEven if training on eval puzzle inputs is allowed, the test input specifically should be forbiddenNo, this is false. The same “transduction” argument applies hereA metalearning benchmark can be transductive in 2 ways:train puzzle $\to$ eval puzzleswithin the eval puzzle, example pair $\to$ test pairThis criticism is specifically answered by the latterThis is against testing policyNo this is false.The policy says “test taker must not know what the test will be”. People interpreted this as saying TTT is banned. But it actually refers to the human designing the AI system, not the AI system itself.Eg: to discourage designing inductive biases based on the eval set.To anyone active in the ARC community, this has always been clear since test time training has been allowed and encouraged. Steven and Chew’s comments clarify this and other concerns.TTT also follows the spirit of a metalearning benchmark, so its fine!You are not including training costsNo, this is false. I show the entire lifetime compute. This is the cost of training the model from init + the total cost of running inference on all tasks. Yes it totally amounts to 67 cents. Check the prices of a 5090 for 2hrs on vast.aiTest time training is traditionally done one task at a time. Training on all test tasks at once is unrealisticYes, this criticism makes sense. But it’s nuancedI agree that its rare to see to face problem sets in real life where every problem is given at once. Even