Mage-Flow: модель на 4 млрд параметров генерирует изображение за 0,59 секунды
Команда во главе с Синьцзе Чжаном представила Mage-Flow, компактную модель генерации и редактирования изображений на 4 млрд параметров. Авторы отмечают, что крупные генеративные модели становятся всё мощнее, но обучение, дообучение и эксплуатация обходятся дорого, и решают эту проблему совместным проектированием двух компонентов стека. Первый, Mage-VAE, лёгкий токенизатор-кодировщик, который сжимает изображение в компактное латентное представление. Он использует одношаговое диффузионное кодирование и декодирование с регуляризацией по опорным латентам (anchor-latent regularization): это сохраняет качество восстановления на уровне сильных открытых VAE-моделей, но снижает вычислительные затраты на токенизацию более чем на порядок. Второй компонент, трансформер-диффузионная модель для генерации в исходном (native) разрешении с мультимодальным вниманием, обучаемая методом rectified flow matching. За счёт упаковки данных в исходном разрешении и слияния CUDA-ядер на уровне всего стека пропускная способность обучения выросла примерно в 2,5 раза, а сама модель научилась работать с изображениями произвольного разрешения без предварительного изменения размера.
На этой основе авторы построили семейство моделей с базовой (Base), дообученной подкреплением (RL-aligned, метод Diffusion-NFT) и ускоренной (Turbo) версиями, для задач как генерации, так и редактирования изображений. Diffusion-NFT улучшает следование инструкциям, качество рендеринга текста на изображении, эстетику и точность редактирования. Turbo-версии получены дистилляцией в несколько шагов с состязательным перцептивным управлением (adversarial perceptual guidance) и генерируют изображение всего за 4 шага диффузии, что даёт низкую задержку при инференсе. По заявлению авторов, несмотря на компактный размер, Mage-Flow и Mage-Flow-Edit показывают конкурентоспособные результаты на стандартных бенчмарках генерации и редактирования изображений. Главный практический результат, версии Turbo делают генерацию и редактирование в высоком разрешении пригодными для интерактивного использования: при разрешении 1024×1024 на одной видеокарте NVIDIA A100 Mage-Flow-Turbo генерирует изображение за 0,59 секунды, а Mage-Flow-Edit-Turbo редактирует изображение за 1,02 секунды, при этом сохраняя небольшой объём потребляемой видеопамяти. Авторы делают вывод, что тщательное совместное проектирование токенизатора, базовой модели и системного стека позволяет добиться сильных результатов в генерации и редактировании изображений высокого разрешения в рамках эффективного семейства моделей на 4 млрд параметров.
Ключевые факты
- Mage-Flow, компактный генеративный стек на 4 млрд параметров для генерации текст-в-изображение и редактирования изображений по инструкциям
- Токенизатор Mage-VAE снижает вычислительные затраты на токенизацию более чем на порядок за счёт одношагового диффузионного кодирования с регуляризацией по опорным латентам
- Упаковка данных в исходном разрешении и слияние CUDA-ядер повышают пропускную способность обучения примерно в 2,5 раза
- Семейство моделей включает базовую версию, дообученную подкреплением (Diffusion-NFT) и ускоренную Turbo-версию с 4-шаговой дистилляцией, отдельно для генерации и для редактирования
- На одной видеокарте NVIDIA A100 Mage-Flow-Turbo генерирует изображение 1024×1024 за 0,59 секунды, а Mage-Flow-Edit-Turbo редактирует изображение за 1,02 секунды
Почему это важно
Работа показывает, что совместное проектирование токенизатора, диффузионной модели и системного стека может дать почти на порядок более дешёвую генерацию изображений без явной потери качества по сравнению с крупными открытыми моделями. Это смещает разговор о генеративных моделях изображений с гонки за размером параметров в сторону инженерной эффективности всего конвейера, от токенизации до инференса.
Кому это важно
Прежде всего это интересно исследователям и инженерам, которые разрабатывают или дообучают модели генерации и редактирования изображений и упираются в стоимость обучения и инференса на GPU. Также это релевантно командам, создающим интерактивные продукты с генерацией изображений в реальном времени, где решает именно задержка ответа модели.
Как это применить
В тексте статьи описаны архитектурные приёмы, одношаговый VAE с регуляризацией по опорным латентам, упаковка в исходном разрешении, слияние CUDA-ядер, дистилляция в несколько шагов с состязательным перцептивным управлением, которые в принципе переносимы на другие диффузионные стеки для ускорения обучения и снижения задержки инференса. Авторы не сообщают в тексте о публикации кода или весов модели, поэтому судить о готовности стека к прямому использованию пока нельзя.
Можно ли доверять
Материал, препринт (карточка статьи на Hugging Face Papers), заявленные цифры по качеству и скорости получены и представлены самими авторами; независимой сторонней проверки результатов в источнике не приводится. На момент публикации у статьи скромный отклик, 23 балла и 1 комментарий на площадке Hugging Face.
Риски и подводные камни
Замеры скорости (0,59 и 1,02 секунды) сделаны на одной конкретной видеокарте NVIDIA A100 и могут не воспроизводиться на другом железе или в других условиях нагрузки. «Конкурентоспособные результаты на стандартных бенчмарках», формулировка авторов; из текста не следует, что модель превосходит более крупные системы, только что она сопоставима с ними при существенно меньшем размере.