Масштабирование токенизаторов изображений с помощью групповой сферической квантизации
В последние годы наблюдается значительный прогресс в области генеративных моделей для изображений и видео, включая такие подходы, как автогрегрессионные модели, модели маскированного языка и методы диффузии. Эти методы превосходят традиционные модели на основе GAN, что связано с их способностью к более высококачественной генерации изображений. Одной из ключевых характеристик многих из этих моделей является использование дискретных латентных представлений изображений, что особенно актуально в контексте моделей, основанных на языковых моделях. Квантизация этих представлений становится критически важной для достижения высокой степени генерации, поскольку токенизированные изображения способствуют улучшению эффективности модели и качеству генерации.
Однако, несмотря на достижения, существуют значительные проблемы в обучении токенизаторов изображений, включая необходимость балансировки между эффективностью сжатия и точностью восстановления. Многие современные токенизаторы по-прежнему зависят от устаревших гиперпараметров, основанных на GAN, что часто приводит к субоптимальным результатам. В этой статье мы представляем новый подход, известный как Групповая Сферическая Квантизация (GSQ), который использует инициализацию кодовой книги на основе сферического распределения и регуляризацию поиска, чтобы улучшить качество восстановления при меньшем количестве итераций обучения.
Связанные работы
Основой для токенизации изображений служит Вариационный Автоэнкодер (VAE), который изначально разрабатывался для сжатия изображений в непрерывное латентное пространство. Однако традиционные VAE часто сдерживаются сильной регуляризацией KL, что ограничивает их применение в качестве токенизаторов изображений. В отличие от этого, VQ-VAE и VQ-GAN стали предпочтительными благодаря эффективному использованию кодовой книги для регуляризации латентного распределения. Эти модели продвинули обучение токенизаторов, внедрив потери восприятия и потери противодействия, что улучшило качество генерируемых изображений.
Модели, такие как FSQ и LFQ, также продемонстрировали успех в области токенизации изображений, но они жестко связывают размер латентного пространства и размер кодовой книги, что делает независимое масштабирование невозможным. В отличие от них, GSQ предлагает более гибкий подход, позволяя оптимизировать как латентные размеры, так и размеры кодовых книг.
Методология
1. Предварительная информация: VQ токенизатор изображений
Токенизатор изображений включает в себя энкодер и декодер. Энкодер сжимает высокоразрешающее входное изображение в непрерывные латентные карты, а декодер восстанавливает изображение из латентного представления. При использовании векторного квантизатора латентное пространство дискретизируется, сопоставляя латентные векторы с индексами в кодовой книге.
2. Простое масштабирование с GSQ
Для достижения более высокого пространственного сжатия требуется увеличение латентной размерности, что может привести к проблемам с высокой размерностью, ограничивая достигнутые коэффициенты сжатия. GSQ решает эту проблему, разбивая каждый латентный вектор на группы. Это позволяет эффективно сжимать данные без ущерба для точности восстановления. Мы инициализируем элементы кодовой книги из сферического равномерного распределения и применяем нормализацию L2 во время поиска.
Эксперименты
1. Оптимизированное обучение для GSQ-VAE
Мы исследовали эффективность предложенных улучшений для GSQ на VAE-токенизаторах, включая влияние конфигураций обучения, вспомогательных потерь и архитектуры модели. В результате наши эксперименты показали, что использование инициализации кодовой книги на основе сферического распределения значительно улучшает использование кодовой книги и снижает потери восстановления.
2. Эффективность сферической квантизации
Анализ показывает, что инициализация кодовой книги из сферического равномерного распределения значительно улучшает использование кодовой книги, достигая близкого к 100% использования в процессе обучения. Это позволяет достичь значительного снижения rFID (размерности восстановления) и улучшения качества генерируемых изображений.
3. Масштабирование латентного пространства и размера кодовой книги
Мы исследовали влияние масштабирования латентной размерности и размера кодовой книги на качество восстановления. Результаты показывают, что увеличение размера кодовой книги в сочетании с меньшими латентными размерами последовательно дает превосходные результаты восстановления, что подтверждает значимость большого словаря кодовой книги для улучшения представительных возможностей квантизатора.
Заключение
Мы представили новый метод квантизации, известный как Групповая Сферическая Квантизация (GSQ), который включает в себя инициализацию кодовой книги на основе сферического распределения и декомпозицию латентных векторов. Наши эксперименты показывают, что GSQ эффективно масштабируется в высокоразмерных латентных пространствах, позволяя достичь высококачественного восстановления с меньшим количеством итераций обучения. Эти результаты открывают новые возможности для более эффективного и масштабируемого обучения токенизаторов изображений, что в свою очередь способствует улучшению генеративных моделей для задач высокой четкости генерации изображений.