Маленький трансформер, большой результат

Митхил Вакде обучил небольшой трансформер с нуля за 1.5 часа на видеокарте RTX 5090, и модель обошла многие крупные LLM, сравнявшись по результатам с TRM/HRM. Итоговый счёт — 44% на публичном eval ARC-AGI-1 при стоимости прогона около $0.67. Предыдущий результат того же автора многие считали невозможным — он привлёк внимание таких исследователей, как Lucas Beyer, Jeremy Howard и Rohan Anil, и разошёлся вирусно в X.

ℹ Что такое ARC-AGI?
ARC-AGI (Abstraction and Reasoning Corpus) — бенчмарк для оценки абстрактного мышления ИИ. Каждое задание состоит из 2–5 демонстрационных пар «вход → выход», по которым нужно вывести правило и применить его к новым входным данным. Бенчмарк использует цветные сетки до 30×30 клеток с 10 цветами и спроектирован для измерения человекоподобного обобщённого интеллекта при минимуме готовых знаний.

Почему это важно

Автор считает sample efficiency (эффективность обучения на малом числе примеров) важнейшей проблемой современного ИИ. Цель работы — найти пределы sample efficiency в рамках трансформеров и снизить стоимость итераций.

ARC-AGI идеально подходит для этого: всего 1000 головоломок, каждая с уникальным правилом, и при этом бенчмарк до сих пор далёк от насыщения. Автор признаётся, что сам не ожидал достичь 45% с одним лишь трансформером — он думал, что для этого понадобятся принципиально новые идеи.

«Я не понимаю, почему другие не додумались до этого раньше. Это просто трансформер с самым очевидным представлением данных. Бенчмарк открыт 6 лет, пользовался широкой известностью и имел приз в миллион долларов!»

Как это работает

Подход строится на test-time training (обучении во время теста): модель обучается с нуля на каждом конкретном наборе задач прямо на этапе инференса.


graph TD
    A[Входные пары вход→выход] --> B[Токенизация последовательностей]
    B --> C[Авторегрессионное обучение трансформера]
    C --> D[3D RoPE позиционные эмбеддинги]
    D --> E[Per-task аддитивный эмбеддинг]
    E --> F[Аугментации цвета и диэдральных перестановок]
    F --> G[Инференс: топ-2 варианта ответа]
    G --> H[44% на ARC-AGI-1]

Ключевые улучшения по сравнению с предыдущей версией

ИзменениеЭффект
SwiGLU вместо GELU, RMSNorm вместо LayerNormРост точности
8 слоёв вместо 4Рост точности
AdamW → NorMuonСтабильная сходимость, меньше затрат
Flash Attention + Flex AttentionУскорение обучения
Только выходные токены в лосс-функции+4 п.п. (40% → 44%)
Меньше аугментацийСнижение стоимости
Данные из ARC-2 (без утечки)Дополнительный прирост
💡 NorMuon вместо AdamW
NorMuon (Neuron-wise Normalized Muon) — оптимизатор, сочетающий ортогонализацию обновлений и адаптивные learning rate на уровне отдельных нейронов. Эксперименты показали, что NorMuon превосходит Adam на 21.74% по эффективности обучения и Muon на 11.31% на модели 1.1B параметров, сохраняя при этом сопоставимый объём памяти.

Что дают аблации

Самый важный вклад — качество представлений:

  • Удаление 3D RoPE (трёхмерного позиционного кодирования) или per-task эмбеддинга → результат падает до ~25%
  • Переход от 3D RoPE к 1D → ~24%
  • Обучение только на ARC-1 + ConceptARC → ~40% (примерно то же)
  • Обучение на входных токенах → чуть хуже, ~39%

Интересный эффект

После перехода к обучению только на выходных токенах модель стала supervised (контролируемой). Странность: тестовый loss при этом ухудшился, но итоговый score вырос, а обучение стало стабильнее и менее случайным. Это указывает на потенциальную ловушку для тех, кто оптимизирует исключительно val loss на малых датасетах.

⚠ Важный методологический момент
Автор тщательно исключил из обучающих данных 773 задачи ARC-1, которые пересекаются с ARC-2, чтобы избежать утечки данных. Без этой фильтрации наивное использование ARC-2 дало бы 100% — что было бы читерством.

Что дальше и как поучаствовать

Код полностью open source. Автор считает, что достичь 65% можно без больших модификаций — объединив решения из нескольких прогонов, он уже получил 55%. По его оценке, стоимость обучения можно снизить ещё в 10 раз с помощью оптимизированного GPU-кода.

Направления для улучшений:

  • Заменить RoPE на PoPE или новый тип позиционного кодирования
  • Дальнейшая модернизация архитектуры
  • Отказ от аугментаций данных (сохраняя низкую стоимость)

Контекст: состояние бенчмарка

ARC Prize 2025 продемонстрировал устойчивый прогресс open source в направлении AGI: лучший результат на Kaggle по ARC-AGI-2 достиг 24%, а число поданных работ выросло до 90 против 47 в 2024 году. На этом фоне 44% на ARC-AGI-1 от маленького трансформера стоимостью менее доллара выглядит тем более впечатляюще.

ARC-AGI-3, первое крупное изменение формата с 2019 года, ставит задачи интерактивного рассуждения и требует от ИИ способностей к исследованию, планированию, памяти и постановке целей. Подходы с высокой sample efficiency, подобные этому, могут оказаться ключевыми для работы с новым поколением бенчмарков.