Samsung TRM on ARC-AGI Benchmark Evaluation

# Testing Samsung's Tiny Recursive Model (TRM) on ARC-AGI-1 **Reproduce → Visualize → Critique → Verdict** Samsung SAIL Montréal's **Tiny Recursive Model (TRM)** is a **7M-parameter** recursive reasoning network that reportedly reaches **45% on ARC-AGI-1** (best config) and **40% Pass@1** for the public [`arcprize/trm_arc_prize_verification`](https://huggingface.co/arcprize/trm_arc_prize_verification) checkpoint — beating many frontier LLMs at <0.01% of their size. This notebook does four things: 1. **Reproduce** the ARC-AGI-1 eval using the official repo + the ARC Prize verification checkpoint (target ≈ 40% Pass@1). 2. **Visualize** a few correctly-solved puzzles as colored ARC grids (input / ground-truth / prediction). 3. **Stress-test the central critique** with three ablations: - **Remove puzzle IDs** — replace the learned puzzle-ID embedding with a blank/random token. - **Reduce recursion** — decode the answer at recursion step 1 instead of the full depth. - **Turn off augmentation voting** — a single canonical forward pass, no 1000-sample majority vote. 4. **Verdict** — does the score reflect *genuine reasoning*, or mostly *benchmark-specific tricks* (puzzle-ID conditioning + test-time compute)? ### What prior evidence predicts A dedicated critique of *this exact checkpoint* (Roye-Azar et al., arXiv:2512.11847) found: | Ablation | Reported effect on ARC-AGI-1 Pass@1 | |---|---| | Baseline (1000-aug + voting) | **40.00%** | | No augmentation voting (singl…

import subprocess, sys, os

# --- Inference-only dependencies (skip training-only adam-atan2 / torchrun) ---
# Sandbox already has torch 2.3.1+cu121 — DO NOT reinstall torch.
DEPS = ["einops", "pydantic", "omegaconf", "hydra-core", "coolname", "argdantic", "numba", "huggingface_hub"]
subprocess.run([sys.executable, "-m", "pip", "install", "-q", *DEPS], check=True)
print("deps installed")

REPO = "/home/user/TinyRecursiveModels"
assert os.path.isdir(REPO), "repo not cloned — run the clone step first"
if REPO not in sys.path:
    sys.path.insert(0, REPO)

import torch
print("torch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
assert torch.cuda.is_available(), "No GPU — TRM eval needs CUDA"
print("device:", torch.cuda.get_device_name(0))
…

deps installed
torch: 2.3.1
CUDA available: True
device: NVIDIA A100 80GB PCIe

checkpoint dir: /home/user/trm_ckpt/arc_v1_public
contents: ['all_config.yaml', 'losses.py', 'step_518071', 'trm.py']
checkpoint file size (MB): 1822.2

import subprocess, sys, os, json, time

REPO = "/home/user/TinyRecursiveModels"
DATA_DIR = os.path.join(REPO, "data", "arc1concept-aug-1000")

# Official builder: augments ARC-AGI-1 (training+evaluation+concept) with num_aug=1000
# and writes the deterministic puzzle-identifier map the checkpoint was trained against.
if os.path.exists(os.path.join(DATA_DIR, "identifiers.json")):
    print("dataset already built, skipping")
else:
    t0 = time.time()
    proc = subprocess.run(
        [sys.executable, "-m", "dataset.build_arc_dataset",
         "--input-file-prefix", "kaggle/combined/arc-agi",
         "--output-dir", "data/arc1concept-aug-1000",
         "--subsets", "training", "evaluation", "concept",
         "--test-set-name", "evaluation"],
        cwd=REPO, capture_output=True, text=True,
…

ation not full, only 575
[Puzzle e5c44e8f] augmentation not full, only 576
[Puzzle a3f84088] augmentation not full, only 576
[Puzzle c97c0139] augmentation not full, only 576
[Puzzle 69889d6e] augmentation not full, only 576
[Puzzle b1fc8b8e] augmentation not full, only 72
[Puzzle 55059096] augmentation not full, only 576
[Puzzle 1990f7a8] augmentation not full, only 72
[Puzzle 4852f2fa] augmentation not full, only 576
[Puzzle 9772c176] augmentation not full, only 576
[Puzzle e872b94a] augmentat…

import sys, torch
import torch.nn as nn

# The kernel imported an older typing_extensions before the upgrade; purge cached
# modules so the fresh (Sentinel-bearing) version + pydantic reload cleanly.
for _m in [m for m in list(sys.modules) if m.split(".")[0] in
           {"typing_extensions", "pydantic", "pydantic_core",
            "models", "argdantic", "omegaconf", "hydra"}]:
    del sys.modules[_m]

# torch.nn.Buffer was added in torch 2.5; sandbox has 2.3.1. Shim it as a
# non-trainable Parameter so the H_init/L_init attributes register on the module
# (and move to CUDA with .cuda()) and load correctly from the checkpoint.
if not hasattr(nn, "Buffer"):
    def _Buffer(data, persistent=True):
        return nn.Parameter(data, requires_grad=False)
    nn.Buffer = _Buffer

…

missing keys: ['inner.puzzle_emb.local_weights', 'inner.puzzle_emb.local_ids', 'inner.rotary_emb.cos_cached', 'inner.rotary_emb.sin_cached']
unexpected keys: []

Total params:            455,667,715
  of which puzzle-emb:   448,719,872  (98.5%)
  'tiny network' (rest): 6,947,843  (~the advertised 7M)

puzzle_emb table shape:  (876406, 512)
GPU mem allocated (GB): 8.12

# Sanity: confirm the puzzle-embedding table holds trained values (not re-init noise),
# and that rotary caches populate on first use. This matters because the whole critique
# hinges on these 448M learned puzzle-specific parameters.
w = model.inner.puzzle_emb.weights
print("puzzle_emb dtype:", w.dtype, "| device:", w.device)
print("puzzle_emb stats  mean=%.4f std=%.4f min=%.3f max=%.3f" % (
    w.float().mean().item(), w.float().std().item(), w.float().min().item(), w.float().max().item()))
print("row 0 (<blank>) norm: %.4f" % w[0].float().norm().item())
print("row 1 (first real)  norm: %.4f" % w[1].float().norm().item())
# A trained table has non-trivial, heterogeneous row norms; a fresh init would be ~uniform.
row_norms = w[1:5000].float().norm(dim=1)
print("row-norm spread (first 5k real IDs): mean=%.3f std=%.3f" % (row_norms.mean(), row_norms.std()))

puzzle_emb dtype: torch.float32 | device: cuda:0
puzzle_emb stats  mean=0.0013 std=0.1496 min=-1.969 max=1.650
row 0 (<blank>) norm: 0.0000
row 1 (first real)  norm: 5.5002
row-norm spread (first 5k real IDs): mean=3.170 std=0.360

import numpy as np, os, json

DATA_DIR = "/home/user/TinyRecursiveModels/data/arc1concept-aug-1000"
TEST_DIR = os.path.join(DATA_DIR, "test")
print("test files:", sorted(os.listdir(TEST_DIR)))

# The builder writes set-prefixed arrays: all__inputs.npy, all__labels.npy, etc.
def load_set(prefix):
    d = {}
    for name in ["inputs", "labels", "puzzle_identifiers", "puzzle_indices", "group_indices"]:
        p = os.path.join(TEST_DIR, f"{prefix}__{name}.npy")
        if os.path.exists(p):
            d[name] = np.load(p, mmap_mode="r")
    return d

test = load_set("all")
for k, v in test.items():
    print(f"  {k:22s} shape={v.shape} dtype={v.dtype}")
…

test files: ['all__group_indices.npy', 'all__inputs.npy', 'all__labels.npy', 'all__puzzle_identifiers.npy', 'all__puzzle_indices.npy', 'dataset.json']
  inputs                 shape=(385815, 900) dtype=uint8
  labels                 shape=(385815, 900) dtype=uint8
  puzzle_identifiers     shape=(368150,) dtype=int32
  puzzle_indices         shape=(368151,) dtype=int32
  group_indices          shape=(401,) dtype=int32

total augmented rows: 385,815
total puzzles (groups): 400
avg augmentations/pu…

import json, time, sys
import numpy as np
import torch
from tqdm.auto import tqdm
from dataset.build_arc_dataset import inverse_aug, grid_hash, arc_grid_to_np
from evaluators.arc import _crop
from puzzle_dataset import PuzzleDataset, PuzzleDatasetConfig

DATA_DIR = "/home/user/TinyRecursiveModels/data/arc1concept-aug-1000"
SEP = "|||"
BLANK_ID = 0

with open(f"{DATA_DIR}/identifiers.json") as f: IDENTIFIERS = json.load(f)
with open(f"{DATA_DIR}/test_puzzles.json") as f: TEST_PUZZLES = json.load(f)
NUM_IDS = len(IDENTIFIERS)

# Precompute orig-name per identifier index (avoids per-row .split() in the hot loop)
ORIG_OF = np.array([IDENTIFIERS[i].split(SEP)[0] if i != BLANK_ID else "" for i in range(NUM_IDS)], dtype=object)
…