ZipAR: Ускорение автогрессивной генерации изображений через пространственную локальность

Современные достижения в области больших языковых моделей (LLM) с использованием парадигмы "предсказание следующего токена" продемонстрировали выдающиеся возможности в решении текстовых задач. На основе этих успехов многие исследования начали расширять эту парадигму для генерации визуального контента, что привело к разработке автогрессивных (AR) моделей генерации изображений. Эти модели не только создают изображения и видео высокого качества, которые сопоставимы или даже превосходят производительность современных диффузионных моделей, но и способствуют унифицированному многомодальному пониманию и генерации.

Тем не менее, медленная скорость генерации остается значительным барьером для широкого применения этих моделей. Для генерации изображений или видео высокого разрешения модели должны последовательно производить тысячи визуальных токенов, что требует многочисленных проходов вперед и приводит к высокой задержке. В ответ на это возникли различные подходы, направленные на уменьшение количества проходов вперед, необходимых для генерации длинных ответов.

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

Основные концепции ZipAR

Пространственная локальность

Изображения обладают характерной локальной структурой, и пространственно удаленные регионы, как правило, имеют минимальную взаимозависимость. Это наблюдение стало основой для разработки ZipAR. Исследования показывают, что между пространственно соседними токенами существует значительная корреляция, что позволяет нам декодировать токены из следующей строки без ожидания завершения декодирования текущей строки.

Параллельное декодирование

В традиционных автогрессивных моделях визуальной генерации токены генерируются в растре, что означает, что первый токен в строке не может быть сгенерирован, пока не будет декодирован последний токен в предыдущей строке. ZipAR изменяет этот подход, позволяя декодировать токены из разных строк параллельно. Это достигается путем определения фиксированного размера окна, чтобы определить, какие токены являются пространственно смежными.

Алгоритм ZipAR

Алгоритм ZipAR позволяет начать генерацию токенов в следующей строке, как только количество сгенерированных токенов в текущей строке превышает установленный размер окна. Таким образом, несколько строк могут декодироваться одновременно, что значительно увеличивает скорость генерации. Все токены, сгенерированные параллельно, производятся с использованием оригинальной модели без необходимости в дополнительной оценке или обновлениях.

Связанные работы

Автогрессивная генерация изображений

Успехи трансформерных моделей в текстовых задачах вдохновили исследования, направленные на применение автогрессивного моделирования к визуальному контенту. Эти методы можно классифицировать на две основные категории: подходы типа GPT, использующие парадигму предсказания следующего токена, и подходы типа BERT, применяющие модели с маскированным предсказанием. Однако модели, обученные с использованием предсказания следующего токена, как правило, менее эффективны при выборке по сравнению с моделями, которые предсказывают несколько токенов за один проход вперед.

Эффективное декодирование LLM

Существуют два основных подхода к уменьшению количества проходов вперед, необходимых для генерации длинных ответов в LLM. Первый подход включает выборку нескольких кандидатных токенов перед их проверкой с помощью основной LLM. Второй подход, называемый декодированием Якоби, включает случайное предположение о следующих токенах в последовательности, которые затем итеративно обновляются. Однако ни один из этих подходов не был специально разработан для автогрессивной визуальной генерации.

Методология ZipAR

Предварительные сведения

Автогрессивные модели визуальной генерации с парадигмой предсказания следующего токена демонстрируют исключительную универсальность в различных задачах, связанных с взаимодействием между текстом и изображением. В процессе генерации изображений AR модели последовательно производят визуальные токены в растре, а затем полная последовательность токенов перерабатывается через визуальный декодер для восстановления изображений.

Инференс с ZipAR

ZipAR предлагает простой и эффективный параллельный декодирующий механизм для автогрессивных моделей визуальной генерации. Ключевая идея заключается в том, что нет необходимости ждать завершения генерации всей строки, прежде чем начинать декодирование следующей строки. Это позволяет значительно увеличить скорость генерации, не теряя при этом качество выходных изображений.

Эксперименты

Подробности реализации

Для оценки эффективности ZipAR мы интегрировали его с тремя современными моделями автогрессивной визуальной генерации: LlamaGen, Lumina-mGPT и Emu3-Gen. Мы провели эксперименты по генерации изображений с текстовыми подсказками и по условной генерации изображений для различных наборов данных.

Количественные результаты

Результаты показали, что ZipAR может значительно уменьшить количество проходов вперед, необходимых для генерации изображений, при этом минимально влияя на качество изображений. Например, ZipAR с размером окна 16 смог сократить количество проходов на 30,5% при незначительном увеличении FID.

Качественные визуализации

Визуализации результатов генерации изображений показали, что ZipAR может уменьшить количество шагов вперед до 91% при сохранении высокого качества изображений, богатых семантической информацией.

Заключение и будущее

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