Meta AI представляет Token-Shuffle: простой метод уменьшения количества токенов изображений в трансформерах

Проблемы автокорреляционной генерации изображений

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

Введение Token-Shuffle

Meta AI разработала Token-Shuffle — метод, уменьшающий количество токенов изображений, обрабатываемых трансформерами, без потери способности предсказывать следующий токен. Метод основан на выявлении избыточности размерности в визуальных словарях мультимодальных больших языковых моделей (MLLM). Визуальные токены, получаемые с помощью векторной квантизации, занимают высокоразмерные пространства, но содержат менее плотную информацию по сравнению с текстовыми токенами.

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

Принцип работы Token-Shuffle

Метод включает две операции: token-shuffle и token-unshuffle. При подготовке входных данных соседние токены объединяются с помощью MLP в сжатый токен, сохраняющий важную локальную информацию. Для размера окна смешивания s количество токенов сокращается в s² раз, что значительно уменьшает количество операций трансформера (FLOPs). После трансформера операция token-unshuffle восстанавливает исходное пространственное расположение с помощью легких MLP.

Этот подход позволяет генерировать изображения с разрешением до 2048×2048 без изменений в архитектуре трансформеров или добавления дополнительных функций потерь и энкодеров.

Адаптированный планировщик управления

Token-Shuffle использует специализированный планировщик classifier-free guidance (CFG) для автокорреляционной генерации. Вместо фиксированного масштаба управления сила управления постепенно настраивается, что снижает артефакты на ранних токенах и улучшает соответствие текста и изображения.

Результаты и оценка

На бенчмарках GenAI-Bench и GenEval Token-Shuffle превзошла другие AR модели и диффузионные методы. На GenAI-Bench с моделью LLaMA на 2,7 млрд параметров достигнут VQAScore 0,77 на сложных запросах, что выше конкурентов на значительные значения. Человеческая оценка подтвердила лучшее соответствие текста, меньшее количество визуальных дефектов и более высокое качество изображений, хотя наблюдалось небольшое уменьшение логической согласованности по сравнению с диффузионными моделями.

Качество изображения и компромиссы

Token-Shuffle генерирует детализированные и согласованные изображения разрешением 1024×1024 и 2048×2048. Исследования показали, что меньшие окна смешивания (например, 2×2) обеспечивают оптимальный баланс между вычислительной эффективностью и качеством, тогда как большие окна дают ускорение ценой потери мелких деталей.

Значение для будущих технологий синтеза изображений

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