Explorative Modeling: новая парадигма обучения ускорила обучение генерации изображений в 6,2 раза

Классическая проблема генеративных моделей: когда у запроса много правильных ответов (например, «сгенерируй собаку», таких изображений миллиарды), модель, обученная напрямую предсказывать ответ, учится выдавать усреднение по всем вариантам, размытое пятно, непохожее ни на одну реальную собаку. Индустрия решает это «факторизацией генерации»: разбивает выдачу на много маленьких шагов, на каждом из которых правильный ответ практически один. Так работают авторегрессионные языковые модели (предсказание по одному токену/фрагменту) и диффузионные модели (сотни маленьких шагов от шума к данным). Но у такого разбиения есть цена: модель обучают на одном шаге, а на инференсе она работает сотни или тысячи шагов подряд, собственные несовершенные предсказания подаются обратно на вход, ошибки накапливаются (эффект exposure bias), поэтому видео у диффузионных моделей «расползается» после десяти секунд, а длинные тексты у языковых моделей теряют связность. Вдобавок такие модели никогда не работают «от начала до конца» одним и тем же способом на обучении и на инференсе.

Авторы предлагают факторизовать не генерацию, а обучение: на каждом шаге модель делает K независимых попыток сгенерировать выход, ошибка считается для каждой, и градиент идёт только через попытку с наименьшей ошибкой, остальные K−1 отбрасываются. Метод назван Explorative Modeling (XM), а теоретическая рамка вокруг него, Mode Forcing: идея, что и факторизация генерации, и обучение через K попыток, это два разных способа сделать так, чтобы оптимальным предсказанием модели становился реальный пример данных, а не их среднее. Авторы вводят термин «генеративная выразительность» (generative expressivity), количество разных ответов, которые модель способна выдать. По их словам, десятилетиями масштабировали число параметров и объём данных, а генеративная выразительность оставалась зафиксированной уже на этапе выбора архитектуры; обучение через K попыток превращает её в третью ось масштабирования наравне с параметрами и данными.

На практике добавление exploration поверх модели RAE (методики, близкой к state-of-the-art в генерации изображений на ImageNet) без изменения остальных гиперпараметров достигает финального качества RAE при 6,2× меньшем объёме данных и 4,1× меньшем числе FLOP, а по скорости сходимости обгоняет RAE в 6,2 раза; итоговая модель (XRAE) достигает FID 1.43 на ImageNet 256 без guidance, около state-of-the-art показателя. На другом бэйзлайне, настроенной модели SiT, обученной на трети обычного вычислительного бюджета, добавление exploration повышает эффективность по FLOP до 52% и достигает того же качества при 2,5× меньшем объёме данных.

Сводные показатели по всем экспериментам: модели с exploration (XM) в целом достигают 6,2× эффективности по данным, 4,1× по вычислениям (FLOP) и на 47% лучшей эффективности по параметрам относительно баз для сравнения; при этом выигрыш растёт вместе с масштабом, с 7% до 36% при росте объёма данных и с 13% до 23% при росте числа параметров. Отдельно, как самостоятельный (end-to-end) подход без диффузионной пошаговости, XM на задачах управления достигает качества диффузионных моделей при вычислениях на инференсе, меньших до 256 раз. Авторы также сообщают, что тестировали добавление exploration на генерации видео и на языковых моделях с маскированной диффузией, но конкретные цифры по этим доменам в доступном фрагменте текста обрываются и не приводятся.

Механика метода на обучении сводится к нескольким строкам кода: цикл из K итераций генерирует по одному образцу, для каждого считается функция потерь относительно реальных данных, а обратное распространение идёт только через минимальную из K потерь. Инференс при этом не меняется: для end-to-end моделей XM генерация остаётся тем же самым одним шагом и на обучении, и на инференсе, без рассинхронизации, свойственной пошаговым диффузионным и авторегрессионным моделям. Код метода выложен на GitHub (alexiglad/XM), у проекта есть отдельная страница explorative-modeling.github.io, и авторы анонсируют отдельную статью с изложением теории Mode Forcing.

Ключевые факты

Explorative Modeling (XM): на каждом шаге обучения модель делает K независимых попыток сгенерировать выход, а градиент идёт только через попытку с наименьшей ошибкой относительно реальных данных
Добавление exploration к модели RAE на ImageNet достигает финального качества RAE при 6,2× меньшем объёме данных и 4,1× меньшем числе FLOP; итоговый FID, 1.43 на ImageNet 256 без guidance, около state-of-the-art
На бэйзлайне SiT, обученном на трети обычного бюджета вычислений, exploration повышает эффективность по FLOP до 52% и достигает того же качества при 2,5× меньшем объёме данных
Сводно по экспериментам: 6,2× эффективность по данным, 4,1× по FLOP, на 47% лучше по параметрам; как самостоятельный end-to-end подход, до 256× меньше вычислений на инференсе при сопоставлении качества с диффузионными моделями на задачах управления
Выигрыш растёт с масштабом: с 7% до 36% при увеличении объёма данных и с 13% до 23% при увеличении числа параметров

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

Все современные генеративные модели, языковые, диффузионные, видео, годами масштабировали два параметра: число весов модели и объём обучающих данных. Третий параметр, введённый авторами термин «генеративная выразительность» (сколько разных ответов модель способна выдать), десятилетиями оставался зафиксированным уже на этапе выбора архитектуры и способа факторизации генерации. Explorative Modeling впервые предлагает превратить его в управляемую и масштабируемую величину: чем больше K попыток даётся модели на обучении, тем больше разных валидных ответов она способна научиться различать, вместо того чтобы усреднять их в размытое пятно.

Кому это важно

В первую очередь, исследователям и инженерам, которые обучают генеративные модели изображений, видео и языка: метод добавляется поверх существующих рецептов обучения (в статье, поверх RAE и SiT) без изменения остальных гиперпараметров и без изменения инференса. Также значим для тех, кто упирается в стоимость обучения больших генеративных моделей: заявленная экономия данных и вычислений в разы напрямую снижает бюджет на достижение того же качества.

Как это применить

Изменение затрагивает только цикл обучения: вместо одной генерации на шаге модель делает K генераций, для каждой считается функция потерь к реальным данным, и обратное распространение идёт только через генерацию с минимальной потерей, это несколько строк кода поверх существующего тренировочного цикла. Инференс не меняется вовсе, что делает метод дешёвым для интеграции в уже существующие пайплайны. Код опубликован на GitHub (alexiglad/XM), у проекта есть отдельная страница с материалами.

Можно ли доверять

Материал, блог-пост независимого автора, опубликованный на личном сайте и вынесенный на обсуждение в Hacker News, а не рецензируемая статья в конференционном сборнике; имя автора и его аффилиация нигде в тексте не указаны. Приведённые числа, количественные бенчмарки против общеизвестных базовых моделей и метрик (RAE, SiT, FID на ImageNet), что даёт возможность проверки, но они являются самоотчётом авторов без независимого воспроизведения. Авторы анонсируют отдельную статью с полным изложением теории (Mode Forcing), которая пока не опубликована; доступный фрагмент текста также обрывается до того, как приведены конкретные цифры по видео и языковым моделям.

Риски и подводные камни

Обучение с K попытками на каждом шаге требует K прямых проходов и расчётов потерь вместо одного, это дополнительная вычислительная нагрузка именно на этапе обучения, даже если инференс не меняется; в тексте отмечается, что оптимальное значение K растёт по ходу обучения, то есть эту нагрузку придётся подбирать динамически. Заявленные результаты по видео-генерации и языковым моделям с маскированной диффузией в доступном фрагменте текста не раскрыты числами, насколько эффект переносится на эти домены, по опубликованному материалу судить нельзя. Независимая рецензируемая публикация с полным описанием теории Mode Forcing пока не вышла.