Спекулятивное декодирование для непрерывной автопрогрессивной генерации изображений

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

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

Спекулятивное декодирование, применяемое в LLM (Large Language Models), стремится уменьшить затраты на инференцию через метод "черновик и проверка". Оно включает в себя меньшую модель черновика, которая генерирует последовательность черновых токенов, которые затем проверяются более точной целевой моделью для принятия или отклонения. Однако, спекулятивное декодирование не было адаптировано для непрерывно-значимых автопрогрессивных моделей визуальной генерации.

Концепции и методы

Непрерывное спекулятивное декодирование

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

Принципы спекулятивного декодирования

• Модель черновика (Draft Model): Генерирует черновые токены, которые могут быть несовершенными, но быстрыми в генерации.
• Целевая модель (Target Model): Проверяет каждый черновой токен и решает, принять его или нет, или же генерировать новый.

Принятие и отклонение токенов

Для непрерывного пространства, мы используем функцию плотности вероятности (PDF) для определения критерия принятия. Если отношение PDF целевой модели к PDF модели черновика больше единицы, токен принимается. В противном случае, токен отклоняется с вероятностью, равной 1 минус это отношение, и затем происходит ресэмплирование из модифицированного распределения.

Введение деноизационной траектории выравнивания

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

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

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

Адаптация к непрерывному пространству

При отклонении токена из модели черновика, мы используем метод принятия-отклонения для сэмплирования нового токена из модифицированного распределения, которое не имеет аналитической формы. Мы устанавливаем правильный верхний предел для этого распределения, чтобы избежать сложной интеграции и упростить процесс сэмплирования.

Экспериментальные результаты

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

Заключение

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