Небольшой трансформер, обученный с нуля за 1,5 часа на видеокарте RTX 5090, набирает результат на уровне TRM/HRM и превосходит многие LLM.

Это улучшенная версия предыдущей модели — быстрее, лучше, дешевле, и всё так же с открытым исходным кодом.

Также модель показывает 7% на ARC-2.

Обсуждение в Twitter, код на GitHub

Это третий материал в серии работ по ARC-AGI. Предыдущие: часть 2, часть 1.

Многие сочли предыдущий результат невозможным. Он привлёк внимание ведущих исследователей и стал вирусным в X — например, его обсуждали Lucas Beyer, Jeremy Howard, Rohan Anil и многие другие.

Зачем этим заниматься

Sample efficiency — эффективное обучение на малом количестве примеров — сейчас, пожалуй, главная нерешённая проблема в AI, и цель этой работы — продвинуться в её решении.

Задача ставилась двойная: (1) найти предел эффективности обучения на малых данных при использовании трансформеров и современных методов глубокого обучения, и (2) снизить стоимость экспериментов, чтобы итерации были быстрее и дешевле.

ARC отлично подходит для такой проверки:

• Очень мало примеров (всего тысяча задач) в пространстве высокой размерности
• Это метаобучающий бенчмарк: каждая задача использует своё правило, хотя встречаются и общие концепции
• Нужен минимум априорных знаний: каждое понятие, встречающееся в eval-наборе, присутствует и в train-наборе
• Людям решать такие задачи очень легко, доступно даже исследователям без мощного оборудования
• Бенчмарк ещё не насыщен (если не считать LLM и подходы с огромным количеством синтетических данных или человеческих индуктивных смещений)

Дальше в планах — новые исследовательские идеи, чтобы преодолеть текущие ограничения, при этом стоимость экспериментов останется низкой, чтобы ими мог заниматься кто угодно.

Технические подробности

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

Общий подход похож на предыдущую версию (полное техническое описание здесь), но добавлен ряд улучшений. Краткое резюме метода:

• Каждая пара вход-выход преобразуется в последовательность токенов. На этих последовательностях небольшой трансформер обучается автореверсивно — с нуля, во время теста, как на train-наборе, так и на задачах eval-набора (тестовые метки скрыты).
• Для переноса знаний между задачами каждой задаче присваивается отдельный обучаемый аддитивный эмбеддинг. Поскольку каждая последовательность содержит две 2D-сетки, позиционная информация задаётся через 3D RoPE-эмбеддинги.
• Последовательности дополняются перестановками цветов и диэдральными преобразованиями. Во время инференса тестовые входы аугментируются, а к полученным выходам применяется обратное преобразование. Отправляются 2 самых частых варианта ответа (AAIVR).

Изменения по сравнению с предыдущей версией

Главная цель — найти улучшения архитектуры и алгоритма, повышающие эффективность обучения на малых данных.

Наибольший прирост в качестве дали:

• Современная архитектура (SwiGLU вместо GELU, RMSNorm вместо LayerNorm и так далее)
• Большее разнообразие данных, лучшее перемешивание
• Масштабирование: 8 слоёв вместо 4

Наибольшее снижение стоимости обеспечили:

• Значительно меньше аугментаций (выше эффективность на малых данных!)
• Переход с AdamW на NorMuon
• flash attention с varlen-обучением + flex-attention ядра для инференса

Ключевое изменение — обучение больше не идёт на входных токенах. Функция потерь теперь учитывает только выходные токены, что делает подход supervised. Это дало небольшой прирост — с 40% до 44%, хотя причина не до конца ясна. Возможно, дело в ограниченной ёмкости модели.

Также был увеличен объём тренировочных данных за счёт непересекающихся задач из ARC-2. Это было сделано с большой осторожностью, чтобы исключить утечку данных. Дополнительные данные можно убрать — результат останется на уровне ~40%, но потребуется вдвое больше вычислений.

Контекст: ARC-2 содержит 773 задачи из ARC-1 и 347 новых. Большинство eval-задач ARC-1 повторяются в ARC-2, поэтому наивное обучение на всём ARC-2 приведёт к утечке данных и результату в 100%. Чтобы этого избежать, повторяющиеся 773 задачи были тщательно отфильтрованы (утечки нет!)

Есть и другие изменения, дающие постепенные улучшения в качестве или скорости. Полный список приведён ниже.

Интересное поведение

Поскольку обучение на входах больше не проводится, подход стал supervised. Странность в том, что test loss при этом выше, а результат — лучше! Кроме того, он более стабилен, разброс между запусками меньше.

Многие сегодня работают над эффективностью обучения на малых данных, ориентируясь на минимальный val loss на небольшом датасете. Это разумный подход, но он выявляет характерный failure mode у такой стратегии.

Обучение в unsupervised-стиле, скорее всего, лучше подходит для некоторых сценариев — этот вопрос сейчас изучается отдельно.

До перехода на NorMuon был опробован обычный Muon. Он ощутимо обучался быстрее AdamW, но loss (и результат) под конец обучения "застревал" вместо того, чтобы сходиться. Резкое снижение momentum и/или learning rate в этой точке помогало, но делать такие ручные правки не хотелось. После перехода на NorMuon проблема исчезла.

Абляции

Наибольший вклад в качество вносят хорошие представления данных (3D RoPE + эмбеддинг под задачу).

• Обучение на входах даёт немного худший результат — ~39%
• Ограничение тренировочного набора только ARC-1 + ConceptARC даёт примерно тот же результат: ~40%
• Замена 3D RoPE на 1D снижает результат до ~24%
• Удаление эмбеддингов под задачу снижает результат до ~24%
• Запуск модели в стиле CompressARC (обучение с нуля отдельно на каждой задаче, unsupervised) снижает результат до ~18%
• CompressARC, но supervised — ~15%

Как можно поучаствовать

Код открыт. Его можно модифицировать, чтобы повысить результат или снизить затраты (только не увеличивать объём тренировочных данных).

Достижимая цель — 65%, и для этого не потребуется много изменений. Подтверждение: объединение всех решённых задач из нескольких запусков даёт 55%. Кроме того, немало задач решены "почти полностью". Несколько идей:

• RoPE смешивает позиционную и содержательную информацию, что, вероятно, снижает качество. PoPE должен работать на том же уровне или лучше. Либо можно придумать новый вид позиционного эмбеддинга
• Архитектуру определённо можно ещё модернизировать

Затраты, вероятно, можно снизить в 10 раз за счёт написанного вручную GPU-кода. Также этого можно добиться архитектурными изменениями.

И наконец, стоит разобраться, как убрать аугментации данных (использование аугментаций вызывает у автора внутренний протест — не стоит слушать тех, кто считает это нормальным). Есть несколько очевидных способов, но задача — сохранить низкую стоимость обучения.

Разное

Честно говоря, достижение 45% с одним лишь трансформером было неожиданностью — казалось, что для этого понадобятся принципиально новые идеи. И уж точно не ожидалось достичь этого при таких низких затратах вычислений. Абляции показывают, что удивительно много производительности сохраняется даже без аугментаций и синтетических данных. Теперь есть уверенность, что 65% достижимы в рамках трансформерной архитектуры.

Непонятно, почему до этого не додумались другие — это просто трансформер с самым очевидным представлением данных. Бенчмарк открыт уже 6 лет, привлекал большое внимание и предлагал призовой фонд в миллион долларов! Возможно, исследователи недооценивают глубокое обучение? Или стоимость экспериментов была слишком высокой, чтобы провести абляции как следует? Или все ослеплены LLM и обвязками (harness) вокруг них?

Приложение

Прежняя критика и поддержка подхода от известных исследователей

Предыдущий результат стал вирусным в X, его обсуждали многие опытные исследователи — как за, так и против. Ветки от Jeremy, Lucas, Susan, Andreas, Yoav и многих других. Ниже перечислены все основные аргументы критиков с ответами.

Обучение на eval-задачах — это читерство / «обучение на тесте»

• Нет, это неверно. «Обучение на тесте» конкретно означает обучение на метках тестовых данных. Метки не использовались при обучении.
• Кроме того, ARC — это метаобучающий бенчмарк, и обучение на eval-задачах — именно то, что от участника ожидается.
  • Терминология: у ARC есть train-набор и eval-набор задач. Каждая задача содержит пары примеров и тестовую пару. Пара — это входная сетка + выходная сетка.
  • В ARC меткой является только выходная сетка тестовой пары в eval-задаче.
  • Эти метки не использовались при обучении. Они скрыты, и при желании их можно удалить заранее.

Обучение на входах eval-задач приводит к утечке информации

• Нет, это неверно. Такой подход называется трансдуктивным выводом и изучается ещё со времён Вапника.
• Кроме того, догма игнорировать входы eval-набора не имеет смысла в мире, где пытаются решить задачу continual learning.
• Другие подходы обучают алгоритм метаобучения, а затем применяют его через chain-of-thought или изменение латентов в рекуррентном цикле. Здесь метаобучение происходит напрямую — через изменение весов единственной forward-функции, что принципиально не отличается.
• Примечание: в новом результате с 44% обучение на входах убрано — оно показывает чуть худший результат.

Даже если обучение на входах eval-задач разрешено, конкретно тестовый вход должен быть под запретом

• Нет, это неверно. Здесь применим тот же аргумент о трансдукции.
• Метаобучающий бенчмарк может быть трансдуктивным двумя способами:
  • train-задача → eval-задачи
  • внутри eval-задачи: пара примеров → тестовая пара
• Именно на второй случай отвечает этот аргумент.

Это противоречит правилам тестирования

• Нет, это неверно.
• Формулировка правил гласит: «тестируемый не должен знать содержание теста». Многие интерпретировали это как запрет test-time training. Но речь идёт о человеке, разрабатывающем систему ИИ, а не о самой системе.
  • Например: чтобы не поощрять встраивание индуктивных смещений на основе eval-набора.
• Для всех, кто активно следит за сообществом ARC, это было ясно всегда: test-time training разрешён и даже поощряется. Комментарии Steven и Chew проясняют этот и другие вопросы.
• TTT также соответствует духу метаобучающего бенчмарка — так что всё в порядке!

Затраты на обучение не учтены

• Нет, это неверно. Показана полная стоимость жизненного цикла: стоимость обучения модели с инициализации + суммарная стоимость инференса на всех задачах. Да, в сумме это действительно 67 центов. Можно проверить цену аренды RTX 5090 на 2 часа на vast.ai.

Test-time training традиционно выполняется по одной задаче за раз. Обучение сразу на всех тестовых задачах — нереалистичный сценарий

• Да, в этой критике есть смысл. Но здесь всё не так однозначно.
• Действительно, в жизни редко встречаются наборы задач, где все задачи выдаются одновременно. Даже если это так (например, на экзамене), человек, как правило, решает их по одной.
• Но то, что у людей нет такой способности, не означает, что нельзя строить AI-модель с этой способностью. Иначе можно было бы сказать, что LLM нереалистичны, поскольку человек не может обучиться на всём интернете или читать токены с такой скоростью.
• Кроме того, неясно, действительно ли люди ограничены обработкой одной задачи за раз — возможно, они способны обучаться на разных данных из нескольких сенсорных каналов одновременно.

Указание стоимости на задачу занижает её за счёт амортизации, так как все тестовые задачи обучаются одновременно. Сравнение с другими моделями нечестно

• Да, это справедливое замечание. В защиту можно сказать следующее:
  • Именно так организаторы сравнивают все модели, включая TRM, которая тоже обучается сразу на всех тестовых задачах.
  • Подход был даже более щедрым по отношению к другим моделям — учитывались и обучение, и инференс, тогда как LLM и другие модели не включают затраты на предобучение/офлайн-обучение.
• Теперь используется другой подход: (a) показывается стоимость за весь жизненный цикл, а не на задачу, (b) сравнение проводится только с TRM, HRM и CompressARC, без LLM и других методов, (c) добавлены абляции со сравнимым стилем обучения.

Ответ на критику самого ARC-AGI

После публикации предыдущего результата было много споров вокруг самого ARC-AGI. Часть замечаний обоснована, но многие — вопросы, на которые Chollet уже неоднократно отвечал ранее:

• Что вообще проверяет ARC? (Подвижный интеллект, fluid intelligence)
• Зачем нам заботиться об ARC? (Fluid intelligence до сих пор не решена полностью)
• Решение ARC-AGI не приведёт к AGI (никто такого и не заявлял)
• ARC постоянно двигает планку / специально сконструирован против LLM (оба утверждения неверны)

Хорошие источники — статья Chollet и ^{эти твиты}. Кратко его позиция такова: бенчмарк задуман для проверки fluid intelligence, которую он считает необходимой, но не достаточной для AGI. Решение ARC-1/2 подразумевает ненулевой уровень fluid intelligence, но не является верхней границей. Бенчмарки не сигнализируют о достижении AGI — их цель обозначить верные исследовательские вопросы. Планка не двигалась: ARC-1 появился до LLM, ARC-2 был объявлен до появления ChatGPT, а ARC-3 объявлен до насыщения ARC-2. Chollet рад прогрессу на ARC, поскольку он документирует прогресс в AI в целом.

Главный интерес к ARC связан с тем, что бенчмарк позволяет проверять эффективность обучения на малых данных — важную нерешённую проблему сегодня. Это также хорошо сконструированный метаобучающий бенчмарк, доступный даже тем, у кого нет мощного GPU. Исторически он хорошо выявлял сильные и слабые стороны LLM. Отдельно впечатляет, что бенчмарк оставался ненасыщенным 6 лет, несмотря на высокий профиль и крупный призовой фонд — теперь понятно, что глубокое обучение способно показывать выдающиеся результаты на ARC-1/2.

Некоторые претензии, на взгляд автора, действительно обоснованы:

• Синтетические данные для ARC-1/2 следует запретить
  • Это противоречит духу бенчмарка, однако большинство сегодняшних топовых результатов опираются на большие объёмы синтетики
    • Синтетические данные снижают требуемый уровень fluid intelligence для решения задач
    • Это имело смысл только до 2024 года, когда результаты глубокого обучения были слабыми
    • Теперь известно, что LLM/DL способны выучить что угодно при достаточном объёме тренировочных данных
  • Это также сделало бы бенчмарк хорошим тестом эффективности обучения на малых данных — отличным дополнением к ARC-3
  • Вопрос — как предотвратить использование синтетических данных? Решение простое:
• Запретить офлайн-обучение/предобучение. Модели должны обучаться с нуля после отправки заявки
  • Раньше это считалось невозможным правилом. Эта модель показывает, что это осуществимо
  • Гарантирует отсутствие синтетических данных
  • Делает сравнение моделей честным. Иначе некоторые модели, например LLM, могут «дожимать» бенчмарк за счёт колоссальных объёмов офлайн-обучения (учитывая, что бенчмарк существует давно, накопилось много похожих на ARC датасетов)

• Единый график лидерборда, сравнивающий разные типы моделей, не имеет смысла. Это порождает три проблемы (решение — отдельные графики для разных категорий):
  • Ось X — стоимость за задачу, но она учитывает только онлайн-вычисления. У некоторых моделей (например, LLM) огромные фазы офлайн-предобучения, стоимость которых не учитывается. Бесконечные вычисления на предобучении фактически переносят тестовый набор «в распределение», поэтому такие модели нужно оценивать отдельно.
  • Деление стоимости на количество задач не имеет смысла для моделей, обучающихся сразу на всех тестовых задачах (как эта модель, TRM и HRM)
  • Сравнение LLM на публичном eval-наборе не имеет смысла, поскольку ответы на публичные задачи доступны в интернете
• Организаторы сделали поспешные выводы из результатов TRM и HRM, приписав успех рекуррентным циклам и глубокому supervision. Похоже, это смещение возникает из предположения, что чистое глубокое обучение не способно решить ARC (например, базовые LLM всё ещё слабы на ARC-2). С этим есть основания не согласиться.
• Формулировку политики тестирования можно улучшить, чтобы убрать путаницу (пояснение выше).

Ошибки, которые допускают другие подходы

Предположение, что рекурсия — следующий большой прорыв (например: HRM, TRM, блог Arcprize).
В этом есть определённая привлекательность, но недостаточно абляций, чтобы это доказать. Эта модель показывает, что того же уровня качества можно достичь без рекурсии. Единственное подтверждённое преимущество рекурсии — возможность увеличивать вычисления без увеличения объёма перемещения данных в памяти.

Обманчивый маркетинг HRM/TRM: не нравится, что TRM рекламировалась как модель на 7 млн параметров, хотя в ней обучается O(100M+) весов эмбеддингов. Это вводит в заблуждение, делает модель больше похожей на таблицу поиска и ставит под вопрос, что на самом деле определяет её качество. В лучшем случае её следовало назвать моделью с 7M «активных» весов. То же касается HRM. Ни в одном из случаев это не было упомянуто!

Подходы на базе LLM для ARC больше не демонстрируют новых возможностей:
Наблюдать, как LLM поднимаются по лидерборду ARC, было крайне полезно — об этом ниже, — но из их текущих результатов на ARC-1/ARC-2 вряд ли можно извлечь что-то новое:

• Рост показателей LLM теперь в основном определяется post-training (доказательства в следующем разделе) и, вероятно, зависит от объёма синтетических данных. Модели учатся решать задачи ARC, а не общему абстрактному рассуждению
• Слишком много посторонних факторов, чтобы делать выводы из новых результатов. Сравнение LLM по бенчмаркам в целом плохая практика (например, из-за разного объёма тренировочных данных, направленных на конкретный бенчмарк)
• Для LLM значение имеют только приватные результаты. Показатели на публичном лидерборде бесполезны, поскольку ответы доступны в интернете и используются при обучении.
• Использование обвязок (harness) над LLM для повышения результата малоосмысленно. Всё настоящее post-training волшебство происходит внутри лабораторий-разработчиков фронтир-моделей, и они сами могут строить такие обвязки. Вряд ли continual learning будет решён именно с помощью harness.

Обходные пути вокруг "bitter lesson"
Ранее уже приводились аргументы о том, что синтетические данные и аугментации — это плохая практика. Встраивание индуктивных смещений в модель — тоже плохо. Тот факт, что масштабировать этот бенчмарк без подобных «читов» пока не получается, говорит о том, что впереди ещё есть прорывы. Хотелось бы, чтобы больше людей стремились избегать таких приёмов, противоречащих духу "bitter lesson".

Что можно узнать из результатов LLM на ARC-AGI

LLM уже насытили и v1, и v2 этого бенчмарка. Вот какие выводы можно сделать из их прогресса:

ARC-AGI появился раньше LLM. Изначально они показывали очень слабые результаты, что демонстрировало: предобучение не даёт общих способностей к рассуждению, и LLM могут проваливаться на задачах, крайне простых для людей.

Результат в 75% у O1 от OpenAI стал большой победой для LLM. Он показал, что при достаточном объёме данных LLM способны выучить практически любую задачу в рамках post-training. Судя по всему, это и есть аргумент, который часто приводит Sholto Douglas.

С выходом ARC-2 прогресс всех LLM, включая «размышляющие» модели, оказался сброшен. Это говорит о том, что даже post-training не даёт общих способностей к рассуждению — иначе модель, хорошо решающая ARC-1, автоматически показывала бы хороший результат и на ARC-2.

(По сути, модели учатся решать конкретно задачи ARC, а не общему абстрактному рассуждению, и их результат на задаче зависит от того, насколько хорошо она представлена в тренировочных данных. Кстати, не до конца ясно, существует ли «общее рассуждение» вообще — возможно, люди тоже специализированы.)

С тех пор «размышляющие» LLM демонстрируют устойчивый прогресс на ARC-2. Многие считают, что это говорит об улучшении общих способностей к рассуждению, но при этом не замечают, что базовые модели застряли на однозначных числах процентов. В сочетании с другими данными это говорит о следующем:

• Результаты на ARC-2 определяются post-training (вероятно, в значительной степени объёмом синтетических данных ARC)
• Лаборатории «дожимают» бенчмарк (вероятно, потому что клиентам важны показатели на бенчмарках)
• LLM никак не отличаются высокой эффективностью обучения на малых данных

При этом отношение к LLM остаётся весьма оптимистичным. Тенденции на ARC-2 показывают, что качество продолжит расти вместе с ростом вычислений и данных. Впечатляет и то, насколько снизилась стоимость инференса.

Полный список изменений

Изменения, влияющие на динамику обучения:

1. Оптимизатор изменён с AdamW на NorMuon + вспомогательный AdamW
2. Расписание learning rate изменено с warmup+cosine на WSD-расписание (warmup %, hold, затем линейное затухание до минимума)
3. LayerNorm заменён на RMSNorm
4. FFN изменён с Linear -> GELU -> Linear на gated FFN в стиле SwiGLU (chunk + SiLU gate)
5. Weight decay изменён с «только non-attention linear» на явный weight decay по группам: веса attention, эмбеддинги токенов и эмбеддинги задач/диэдральных преобразований — у каждой группы свой коэффициент
6. Целевая функция обучения изменена с outputs["loss"] (unsupervised-стиль, loss по входу и выходу) на только outputs["output_loss"] (supervised-стиль)
7. Батчинг при обучении изменён с умного bucketing по длине на настоящую случайную батчировку (bucketing сохранён для инференса)
8. Неполные («хвостовые») батчи теперь отбрасываются при обучении (принудительно включён drop_last=True)
9. Формирование датасета теперь поддерживает более широкий набор источников (ARC-1, ARC-2, ConceptARC, опционально отфильтрованные задачи из других датасетов, режимы submission/private), что меняет состав тренировочных данных
10. Аугментация цвета изменена с одной глобальной перестановки на уровне эпохи на аугментацию по кортежам для каждого примера (цвет + диэдральное преобразование)
11. Домен перестановки цветов теперь исключает цвета, встречающиеся только в выходе, вместо «слепых» перестановок 1..9
12. Генерация аугментаций теперь дедуплицирует преобразованные входы через хеширование в пределах задачи (выше уникальное разнообразие примеров)
13. Выбор аугментаций теперь циклируется по эпохам с перемешанным порядком кандидатов (поведение «без повторов» в рамках цикла)
14. Изменения гиперпараметров: оптимизатор/hparams, число эпох, лимит/тип аугментаций, глубина (n_layers) и путь к датасету
15. Добавлен новый dihedral_embedding, суммируется с условной информацией токена (даёт лишь незначительный прирост качества)

Ускорения без изменения динамики обучения:

1. Батчи обучения изменены с padded [B,S] на packed поток токенов с cu_seqlens (без pad-токенов на пути обучения)
2. Attention изменён с padded SDPA-маскирования на packed varlen flash-attention (cu_seqlens), плюс flex-attention ядра для декодирования
3. Диэдральная аугментация перенесена из офлайн-расширения датасета в онлайн-выбор аугментаций во время сборки батча (collate)
4. Разбиение на train-сплит на этапе сборки изменено с ("train","test") на ("train",)

Разное:

1. Поведение при resume изменено: обнаружение смены оптимизатора/hparams теперь может сбрасывать/разогревать расписание, меняя динамику продолженного запуска
2. Шаг scheduler изменён на дробный прогресс эпохи при обучении, вместо чисто пошагового cosine-прогресса