Train Small Transformer in 1.5hrs That Beats Many LLMs

You Don’t Need a Frontier LLM for Hard Reasoning

44% on ARC-AGI-1 public eval. 75M parameters. About 1.5 hours on one RTX 5090. Lifetime train+infer compute around $0.67 as reported by the author. That is not a frontier stack – and it still clears a lot of non-reasoning LLM baselines on this benchmark.

I trained a small transformer in 1.5hrs and it beats many LLMs is a real Pareto spike on sample-efficient metalearning, not a chat demo. Code is open. This is the reproduction path – not another Shakespeare nanoGPT walkthrough.

Quick Context: What Actually Dropped

Mithil Vakde’s write-up and the mdlARC repo describe a modernized vanilla transformer trained entirely at test time on ARC train + eval puzzles (test output labels hidden). No giant pretrain. No program synthesis. Grids → token sequences; next-token prediction with per-task embeddings and 3D RoPE so the input and output planes stay aligned.

HN hit hard – 661 points on the thread. Author’s own framing there: this is not an LLM. Target is sample efficiency on a metalearning benchmark where eval concepts already showed up in train. As of that post it also reports ~7% on ARC-2. The recurring fight is “is TTT cheating?” – labels stay hidden; transduction matches how ARC was designed (Chollet’s ARC paper).

Hands-On: Reproduce the 44% Run

Reference box is a recent NVIDIA GPU. 5090 is what the author timed (CUDA >12.8, ideally >13). Older cards run lower presets; wall clock stretches. No local iron? Rent – as of early 2026, on-demand 5090 rates on public trackers often sit near ~$0.33/hr (varies; check live), so ~2h near the claimed ~$0.67 total.

1. Clone and install

git clone https://github.com/mvakde/mdlARC.git
cd mdlARC
python -m venv .venv && source .venv/bin/activate
pip install torch numpy numba matplotlib
# flash-attn required for the fast path - match CUDA/torch per its docs

Flash-attn is the sharp edge. Wrong wheel → you still train, just crawl, or OOM in odd spots. Confirm torch.cuda.is_available() and the driver/CUDA pair before you burn rental hours.

2. Build the dataset (no leakage)

cd dataset_building_scripts
python download_and_group.py
python build_datasets.py arc1 --add-conceptarc --with-filtered
cd ..

Use --with-filtered. ARC-2 reuses 773 ARC-1 puzzles; skip the filter and you can fake a perfect score. Optional proof step from the repo flow: after JSON is built, delete assets_tmp, assets/solutions.json, and the building scripts so label files are gone.

3. Train + infer in one shot

python run_script.py high

Presets in run_script.py:

• low – 90 epochs, lighter augs (smoke test)
• medium – 240 epochs (balance)
• high – 650 epochs, up to ~300 max augments (published ~44% path)

Loss logging is off by default for speed. Turn it back on in the train loop if you need to see whether the run is stuck. Pipeline: joint training across puzzles (shared per-task embedding setup) then AAIVR at inference – color + dihedral augs on test inputs, generate, inverse-transform, majority vote top outputs.

Start on medium if the card is smaller. Jump to high only after flash-attn and data look clean. Author wall-clock for high on a proper 5090: about 1.5-2 h.

What actually moves the needle, per the published ablations: 3D RoPE instead of plain 1D, a learned per-puzzle embedding added to every token, output-only loss in the latest setup, NorMuon, flash/varlen attention. Strip 3D RoPE or the task embed and scores fall to roughly the mid-20s. Train on inputs too (older unsupervised style) lands near ~39%. CompressARC-style per-task-from-scratch sits around ~15-18% in this frame.

Ever tried explaining a color-grid puzzle to a model that only ever saw chat text? Feels like teaching chess with a dictionary. That mismatch is why this tiny specialist exists at all.

Common Pitfalls That Kill Your Score

Broken CUDA / flash-attn first. Fix the env before you chase hyperparameters.

The catch is the ARC-2 filter – without it the number is invalid, full stop.

Low preset will not print 44%. High is the published config. Joint multi-task training is intentional here; the per-task CompressARC pattern is a different (weaker, in these ablations) bet.

Performance: what 44% actually buys you

As of the author blog + repo: 44% ARC-AGI-1 public eval, ~7% ARC-2, 75M params, 8 layers in the latest, full train+infer lifetime ~$0.67 on the 5090 path. Comparison focus is other TTT-style systems (TRM/HRM class), not API LLMs with pretrain cost zeroed out. On cost-per-point and params-per-point it sits far left.

Ablations put most of the lift on representation: task embeds + 3D RoPE. Optimizer and aug schedule cut dollars and hours; they do not replace those two.

Is a 44% grid specialist “better” than GPT-class models? Only on this axis. Zero world knowledge. Cannot chat. That is the feature, not a bug.

When NOT to Use This

Need open-ended language, tools, or long-context chat? Skip. This is a narrow metalearner for colored grids. Laptop CPU or ancient GPU? Speed claims assume modern CUDA + flash-attn. And the ~67-cent figure is not a free lunch you 100× – returns flatten; the author already flags higher spend plateaus.

Later reading if you want the neighborhood: CompressARC, Modded NanoGPT-style sample-efficiency work, and the original ARC paper linked above.

FAQ

Is training on eval puzzles cheating?

No. Test output labels stay hidden. ARC is a metalearning / few-shot adaptation benchmark; test-time training on visible structure is allowed.

Can I hit similar numbers on a 4090 or Colab A100?

Yes – slower. Smoke on low/medium first. Dial epochs and augs for VRAM. Flash-attn still has to install cleanly or the “fast path” is fiction. A rented A100 can finish a lighter public-style run in a few hours; high preset is a longer sit.

Why does a plain transformer suddenly work when earlier ones struggled?

Turns out the exotic architecture was never the missing piece. Joint training history, 3D RoPE lining up the two grids, and one learned embedding per puzzle do most of the sharing across tasks. Kill positional or task embedding alone and you are back near ~25%. Output-only supervised loss edged the older inputs+outputs setup from ~40% to 44% – author says they do not fully know why; finite capacity is the working guess. Little else needed once representation + TTT were right.

Clone the repo. Rent the cheapest clean 5090 you can find. Run python run_script.py high.

Can you push past 44% without adding data? That is the open move the author left on the table.