Готовые ИИ-генераторы картинок научили дорисовывать любые области за малую долю затрат

Авторегрессионные модели генерируют изображение так же, как языковые модели пишут текст: токен за токеном. Картинку сначала режут на дискретные коды, потом трансформер предсказывает их по очереди, обычно по строкам, слева направо и сверху вниз. Такой порядок называют растровым (raster scan). Исследователи из Stony Brook University и ÉTS Montréal предложили способ превратить уже обученную растровую модель в генератор с произвольным порядком. Он умеет дорисовывать пропуски внутри кадра (inpainting), достраивать картинку за границами (outpainting) и редактировать выбранные области. Переобучать с нуля для этого не нужно. Дополнительная адаптация, по оценке авторов, дешевле обучения «любого порядка» с нуля более чем на порядок.

Почему построчная модель не умеет чинить середину

При растровом обучении модель всегда видит только то, что нарисовано выше и левее. Предсказывать токен, у которого известны соседи справа и снизу, она не училась. Поэтому маска посреди кадра её ставит в тупик.

Существуют два способа обойти ограничение, и оба дороги:

• Обучение на смеси порядков с нуля. Например, RandAR требует 300 эпох, а построчная LlamaGen-XL обходится 100. Бюджет вырастает втрое.
• Дискретная диффузия, к которой относится и MaskGIT. Модель итеративно восстанавливает закрытые маской токены с двунаправленным вниманием. Качество хорошее, но KV-кэш (хранение уже посчитанных промежуточных данных, чтобы не пересчитывать их на каждом шаге), спекулятивное декодирование и серверы вроде vLLM с такой схемой не работают.

Компании с готовым авторегрессионным конвейером и серверами вынуждены строить рядом второй стек. Авторы считают такие траты лишними.

Две детали, без которых переучивание ломается

Идея держится на математике: по правилу цепочки вероятностей совместное распределение токенов можно разложить в любом порядке. Значит, модель, хорошо усвоившая картинки в одном порядке, уже знает почти всё нужное. Не хватает умения пользоваться знанием при другом порядке. Для этого понадобились два компонента, и по отдельности они не работают.

Позиционные эмбеддинги «с заглядыванием вперёд» (future-aware). При растровом порядке модель всегда знает, что рисует следующим. При случайном порядке один и тот же контекст может вести к разным целям. Токен неба и токен лица по одним и тем же видимым данным предсказываются по-разному, и градиенты начинают конфликтовать. В языковой модели XLNet эту проблему решали вторым потоком внимания, который удваивает вычисления в слое и ломает KV-кэш. Здесь решение проще. К входу каждого шага добавляется выученный вектор той позиции, которую надо предсказать следующей. Таблица стартует с нулей, добавка делается только на входном слое. Архитектура остаётся обычным одно-потоковым трансформером.

Отдельно сделан выбор позиции для внимания. В LlamaGen, где используются вращательные позиционные эмбеддинги RoPE, и в RAR, где применены таблицы абсолютных позиций, берётся пространственный индекс токена, а не номер шага. Модель сохраняет всё, что знала о соседстве пикселей.

Постепенное усложнение порядков (curriculum). Если сразу дать случайные перестановки, потери огромны, градиенты шумные, и модель вылетает из удачной точки предобучения, забывая накопленное. Авторы вводят коэффициент, который линейно растёт: сначала почти все примеры идут в растровом порядке, потом всё больше случайных. Каждый пример в батче получает свою перестановку. Простое снижение шага обучения тут не помогает: градиент указывал бы в ту же неудобную сторону. Расписание меняет сам состав обучающих данных, и ранние градиенты согласуются с предобученным решением. Если убрать любой из двух компонентов, обучение расходится или качество заметно падает.

Что получилось и где цена

Метод применили к четырём базовым моделям двух архитектур и двух размеров: LlamaGen-L, LlamaGen-XL, RAR-L и RAR-XL. Результаты:

• На дорисовке внутри кадра и расширении кадра на ImageNet адаптированные модели заметно обходят MaskGIT по FID. Эта метрика сравнивает распределение сгенерированных картинок с настоящими, и чем она ниже, тем лучше.
• С MAR-L, который работает с непрерывными токенами и диффузионной головой, результаты сопоставимы. Ни одной из этих задач модель специально не обучали.
• Одни и те же веса справляются с дорисовкой, расширением и редактированием по областям. Меняется только набор известных токенов на входе.
• Цена есть. На LlamaGen-XL FID в обычной генерации по классам вырос. Inception Score, precision и recall остались близки к исходным. Авторы прямо говорят: это обмен качества на гибкость, а не бесплатное превращение.

Сама архитектура осталась причинной, поэтому KV-кэш и параллельное декодирование работают. Модель можно запускать через vLLM с регистрацией собственной модели. Для дискретной диффузии, по словам авторов, это принципиально недоступно.

Что это даёт инженерам

Кто уже владеет предобученной авторегрессионной моделью картинок и инфраструктурой под неё, получает редактирование без второго стека и без повторного обучения. Дорисовка и расширение кадра становятся режимом уже существующей модели. Ограничения тоже понятны. Проверка ограничена ImageNet и четырьмя базовыми моделями, а обычная генерация немного теряет в качестве. Но направление ясно: полезные способности можно доставать из моделей, которые уже обучены. Код опубликован на GitHub в репозитории Ordered-AR/unlockedadaptation.