🚀 I Trained a Small Transformer in 1.5 Hours That Outperforms Many LLMs

By Mithil Vakde (IIT Bombay '23, Engineering Physics) Based in Bangalore, India | [Twitter] [LinkedIn]

The Headline

I successfully trained a compact transformer from scratch using a 5090 GPU in just 1.5 hours. The result? A model that achieves 44% on ARC-AGI-1 for a total cost of only $\approx 67\text{ cents}$. This model matches the performance of TRM/HRM and represents a significant upgrade over my previous iteration—it is faster, more efficient, cheaper, and remains open source.

Quick Stats:

• ARC-1 Score: 44% (Public Eval)
• ARC-2 Score: 7%
• Cost: \0.67$
• Training Time: 90 minutes

"I believe that sample efficiency is the most critical hurdle in AI today, and my goal is to solve it."

---

🎯 Objectives & The ARC Benchmark

The primary motivations for this project are twofold:

1. To determine the absolute limits of sample efficiency when utilizing current deep learning methods and transformer architectures.
2. To slash iteration costs, making the research cycle significantly faster and more accessible.

Why ARC?

The Abstraction and Reasoning Corpus (ARC) is an ideal testing ground because:

• Sparse Data: It features a high-dimensional space with very few samples (only 1,000 puzzles).
• Meta-learning Nature: Every puzzle follows a unique rule, though they share underlying concepts.
• Minimal Priors: The training set contains all the concepts necessary for the evaluation set.
• Human-Centric: It is easy for humans but remains saturated unsaturated for AI (excluding models that rely on massive synthetic datasets or heavy human-coded inductive biases).

---

🛠️ Technical Approach

The general workflow follows my previous methodology but incorporates several critical upgrades.

The Pipeline

Key Implementation Details:

• Tokenization: Each pair is converted into a sequence of tokens.
• Test-Time Training: The model is trained from scratch at test time on both training and evaluation puzzles (with labels hidden for the latter).
• Cross-Task Learning: To distinguish between puzzles, I utilized a learnt additive embedding for each task.
• Positional Encoding: Since the data consists of two 2D grids, I implemented 3D RoPE (Rotary Positional Embeddings).
• Augmentation: I applied color and dihedral permutations to the sequences. During inference, inputs are augmented, and the resulting outputs are processed via an inverse augmentation.

---

📈 Improvements & Changes

I focused on maximizing sample efficiency while minimizing compute.

⬆️ Performance Boosts

• Modernized Architecture: Swapped GELU for SwiGlu and LayerNorm for RMSnorm.
• Scaling: Increased depth from 4 layers $\to$ 8 layers.
• Data Quality: Improved shuffling and increased data diversity.

⬇️ Cost Reductions

• Efficient Augmentation: Reduced the number of augmentations to improve sample efficiency.
• Optimized Kernels: Utilized AdamW $\to$ NorMuon, Flash Attention with varlen training, and flex attention kernels for inference.

🔄 The Supervised Shift

A major pivot was stopping training on input tokens. The loss function now only considers output tokens.

• Result: Performance climbed from $40\% \to 44\%$.
• Theory: This might be due to finite model capacity.

📦 Integrating ARC-2

I expanded the training set by adding non-overlapping tasks from ARC-2.

• The Leak Problem: ARC-2 contains 773 puzzles from ARC-1. Naive training would cause a data leak, leading to a fake 100% score.
• The Solution: I carefully filtered out those 773 repeated puzzles.
• Impact: While the model still hits $\sim 40\%$ without this data, adding it boosts the score to 44% and reduces the required compute by half.

---

🧐 Observations & Ablations

The Loss Paradox

Interestingly, by switching to a supervised approach (ignoring input tokens), the test loss increased, but the actual score improved. This suggests that minimizing validation loss on small datasets isn't always the best proxy for performance.

Optimizer Insights

While AdamW was fast, the loss tended to plateau. I found that drastically lowering the learning rate or momentum helped, but switching to NorMuon eliminated the problem entirely without manual tuning.

Ablation Study Results

The following table summarizes the impact of various components:

---

🤝 How to Contribute

I want to keep costs low so that researchers globally can participate. I hate that I used high-end GPUs, ignore anyone who thinks it's okay.

The Goal: Reach 65% accuracy. (I've already seen 55% by taking the union of solved tasks across multiple runs).

Suggested Areas for Improvement:

• Positional Embeddings: RoPE mixes content and position; perhaps a new embedding method would work better.
• Architecture: Further modernize the transformer block.
• GPU Optimization: Write handmade GPU code to reduce costs by another 10x.
• Constraint: $\text{Please do not simply increase the training data}$.

Final Thoughts

To be honest, I didn't expect a standard transformer to hit 45%—I assumed we would need entirely new conceptual breakthroughs. I certainly di—