APEX сжимает латенты диффузии в 128 раз и обходит 256-токенную модель на ImageNet-512

Исследователи представили на arXiv метод APEX, который атакует главную проблему дешёвой диффузии: чем сильнее сжат латент, тем труднее его генерировать. На ImageNet-512 APEX превосходит существующие методы при каждом проверенном числе каналов в режимах сжатия 64x и 128x. При равных затратах на обучение и guided-сэмплинг он также обходит модель на 256 токенов на всех проверенных бюджетах, причём чем меньше вычислений, тем заметнее разрыв. Код опубликован на странице проекта.

Дилемма, ради которой всё затевалось

Стоимость латентной диффузии напрямую зависит от числа токенов, с которыми работает денойзер. Меньше токенов означает быстрее и дешевле как обучение, так и инференс. Но информацию из изображения никуда не деть. Если токенов мало, её приходится упаковывать в каналы каждого токена.

Здесь и возникает противоречие, которое авторы называют reconstruction-generation dilemma. Больше каналов даёт автоэнкодеру точнее восстанавливать картинку. Одновременно такой плотный латент сложнее моделировать диффузионной модели. Выигрыш в реконструкции съедается просадкой в генерации.

Что предлагает APEX

Название расшифровывается как Asynchronous Prefix denoising for EXtreme compression. Решение состоит из двух частей, которые работают в паре.

Первая часть — иерархическое латентное пространство. Автоэнкодер упорядочивает каналы по префиксам: сначала идут те, что несут структуру изображения, затем те, что отвечают за детали. Начало вектора каналов описывает общую картину, продолжение её уточняет.

Вторая часть — асинхронный денойзинг по группам каналов. Один общий денойзер генерирует разные группы в разные локальные моменты времени. Число токенов при этом не растёт, а архитектура не усложняется отдельными моделями под каждую группу. По описанию авторов, такая схема позволяет генерировать структуру и детали по-разному, а не сваливать весь латент в одну задачу.

Цифры из статьи

• Бенчмарк: ImageNet-512.
• Режимы сжатия: 64x и 128x.
• Сравнение с существующими методами: APEX лучше при каждом проверенном числе каналов.
• Сравнение с 256-токенной моделью: APEX лучше на каждом проверенном бюджете при равных затратах на обучение и guided-сэмплинг.
• Динамика: чем ниже вычислительный бюджет, тем больше выигрыш.
• Дополнительные токены: не требуются.

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

Что с этим делать на практике

Командам, которые строят собственные генеративные модели изображений, APEX даёт готовый повод пересчитать компромисс между числом токенов и качеством. Код открыт, поэтому воспроизвести результат на ImageNet-512 можно без реверс-инжиниринга статьи.

Осторожность тоже нужна. В аннотации заявлены результаты только на ImageNet-512. Работает ли подход в text-to-image на больших датасетах, с видео или с другими архитектурами денойзера, из неё не видно. Не указана и цена самого автоэнкодера с иерархией каналов: сколько стоит его обучение и насколько он чувствителен к настройкам.

Стоит прочитать полный текст и ответить на три вопроса. Как именно устроено расписание асинхронных времён. Какой выигрыш в реальном времени инференса даёт сжатие 128x. Сохраняется ли преимущество при масштабировании модели. От этих ответов зависит, станет ли APEX рабочим инструментом или останется аккуратной академической работой.