Эффективное декодирование в визуальном авто-регрессионном моделировании: концепция Collaborative Decoding

В последние годы область генерации изображений значительно продвинулась благодаря развитию авто-регрессионных (AR) моделей. Одной из наиболее интересных концепций в этой области является визуальное авто-регрессионное моделирование (VAR), которое предлагает инновационный подход к предсказанию на следующем уровне. Однако, несмотря на свои преимущества, такие как высокая эффективность и масштабируемость, VAR сталкивается с проблемами, связанными с потреблением памяти и вычислительными издержками. В этой статье мы рассмотрим концепцию Collaborative Decoding (CoDe), предложенную для решения этих проблем, а также ее влияние на эффективность и качество генерации изображений.

Введение в визуальное авто-регрессионное моделирование

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

Тем не менее, несмотря на меньшее количество шагов декодирования, подход VAR приводит к значительному увеличению общей длины последовательности. Например, для генерации изображения размером 16x16 токенов, VAR должен декодировать до 680 токенов на 10 масштабах, что в 2.7 раза превышает длину последовательности, необходимую для традиционных AR моделей. Это приводит к значительным затратам памяти, особенно на финальных масштабах, где KV-кеш (кэш ключей и значений) может потреблять до 12 раз больше памяти, чем требуется для прямого вычисления.

Проблемы и ограничения VAR

Основные проблемы, с которыми сталкивается VAR, включают:

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

2. Избыточные вычисления: Поскольку длина последовательности увеличивается, вычислительные затраты на самовнимание растут квадратично, что делает процесс менее эффективным.

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

Концепция Collaborative Decoding (CoDe)

Для решения вышеупомянутых проблем была предложена концепция Collaborative Decoding (CoDe). Эта стратегия направлена на оптимизацию процесса декодирования в рамках VAR, используя два модели различного размера с специализированными ролями.

Основные идеи CoDe

1. Разделение ролей: CoDe использует большую модель VAR в качестве "черновика" для генерации низкочастотного контента на меньших масштабах, в то время как меньшая модель служит "усовершенствователем", сосредотачиваясь на предсказании высокочастотных деталей на больших масштабах.

2. Снижение потребления памяти: Поскольку большая модель используется только для предсказания первых токенов, а затем освобождает KV-кеш, общие затраты памяти значительно снижаются.

3. Специализированная дообучение: Каждая модель проходит специализированное дообучение на своих заданных масштабах, что позволяет избежать взаимных помех между обучением низкочастотных и высокочастотных компонентов.

Эффективность CoDe

Эксперименты показывают, что CoDe обеспечивает значительное улучшение в скорости и снижении потребления памяти. Например, CoDe достигает ускорения в 1.7 раз и снижает использование памяти на 50% при незначительном увеличении FID (Fréchet Inception Distance) с 1.95 до 1.98. При уменьшении числа этапов черновика до 6, CoDe может достичь впечатляющего ускорения в 2.9 раза, что соответствует 41 изображению в секунду при разрешении 256x256 на одном GPU NVIDIA 4090.

Сравнение с традиционными AR моделями

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

Результаты экспериментов

В рамках экспериментов, проведенных для оценки эффективности CoDe, было установлено, что:

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

Заключение

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