Упрощенное декодирование 1: Одношаговая выборка изображений авторегрессионных моделей с соответствием потоков

Distilled Decoding 1: One-step Sampling of Image Auto-regressive Models with Flow Matching

December 22, 2024

Авторы: Enshu Liu, Xuefei Ning, Yu Wang, Zinan Lin

cs.AI

Аннотация

Авторегрессионные (AR) модели достигли передовых показателей в генерации текста и изображений, но страдают от медленной генерации из-за пошагового процесса токен-за-токеном. Мы задаем амбициозный вопрос: можно ли адаптировать предварительно обученную AR модель для генерации выводов всего за один или два шага? В случае успеха это значительно продвинет разработку и применение AR моделей. Мы замечаем, что существующие работы, пытающиеся ускорить генерацию AR путем генерации нескольких токенов одновременно, фундаментально не могут уловить распределение вывода из-за условных зависимостей между токенами, что ограничивает их эффективность для генерации на несколько шагов. Для решения этой проблемы мы предлагаем метод Distilled Decoding (DD), который использует сопоставление потоков для создания детерминированного отображения из гауссовского распределения в распределение вывода предварительно обученной AR модели. Затем мы обучаем сеть для упрощения этого отображения, обеспечивая генерацию на несколько шагов. DD не требует обучающих данных оригинальной AR модели, что делает его более практичным. Мы оцениваем DD на передовых моделях AR для изображений и представляем многообещающие результаты на ImageNet-256. Для VAR, требующего генерации на 10 шагов, DD позволяет генерацию за один шаг (ускорение в 6.3 раза), с приемлемым увеличением FID с 4.19 до 9.96. Для LlamaGen, DD сокращает генерацию с 256 шагов до 1, достигая ускорения в 217.8 раза с сопоставимым увеличением FID с 4.11 до 11.35. В обоих случаях базовые методы полностью терпят неудачу с FID>100. DD также превосходит в генерации текста-изображения, сокращая генерацию с 256 шагов до 2 для LlamaGen с минимальным увеличением FID с 25.70 до 28.95. Как первая работа, демонстрирующая возможность генерации за один шаг для моделей AR изображений, DD бросает вызов преобладающему представлению о том, что AR модели по своей сути медленные, и открывает новые возможности для эффективной генерации AR. Веб-сайт проекта находится по адресу https://imagination-research.github.io/distilled-decoding.

English

Autoregressive (AR) models have achieved state-of-the-art performance in text and image generation but suffer from slow generation due to the token-by-token process. We ask an ambitious question: can a pre-trained AR model be adapted to generate outputs in just one or two steps? If successful, this would significantly advance the development and deployment of AR models. We notice that existing works that try to speed up AR generation by generating multiple tokens at once fundamentally cannot capture the output distribution due to the conditional dependencies between tokens, limiting their effectiveness for few-step generation. To address this, we propose Distilled Decoding (DD), which uses flow matching to create a deterministic mapping from Gaussian distribution to the output distribution of the pre-trained AR model. We then train a network to distill this mapping, enabling few-step generation. DD doesn't need the training data of the original AR model, making it more practical.We evaluate DD on state-of-the-art image AR models and present promising results on ImageNet-256. For VAR, which requires 10-step generation, DD enables one-step generation (6.3times speed-up), with an acceptable increase in FID from 4.19 to 9.96. For LlamaGen, DD reduces generation from 256 steps to 1, achieving an 217.8times speed-up with a comparable FID increase from 4.11 to 11.35. In both cases, baseline methods completely fail with FID>100. DD also excels on text-to-image generation, reducing the generation from 256 steps to 2 for LlamaGen with minimal FID increase from 25.70 to 28.95. As the first work to demonstrate the possibility of one-step generation for image AR models, DD challenges the prevailing notion that AR models are inherently slow, and opens up new opportunities for efficient AR generation. The project website is at https://imagination-research.github.io/distilled-decoding.