Групповая оптимизация самораспространяющихся кодовых книг в моделях векторного квантования

Основная информация

• ID статьи: 2510.13331
• Название: Group-Wise Optimization for Self-Extensible Codebooks in Vector Quantized Models
• Авторы: Hong-Kai Zheng, Piji Li (Нанкинский университет аэронавтики и астронавтики)
• Категория: cs.CV
• Время публикации/Конференция: ICLR 2026
• Ссылка на статью: https://arxiv.org/abs/2510.13331

Аннотация

Vector Quantized Variational Autoencoders (VQ-VAE) осуществляют самообучение с учителем посредством задачи реконструкции, используя ближайший вектор из кодовой книги для представления непрерывного вектора. Однако в моделях VQ по-прежнему существуют проблемы, такие как коллапс кодовой книги. Для решения этих проблем существующие методы используют неявные статические кодовые книги или совместную оптимизацию всей кодовой книги, но эти подходы ограничивают способность обучения кодовой книги, что приводит к снижению качества реконструкции. В данной работе предлагается Group-VQ, метод групповой оптимизации кодовой книги. Каждая группа оптимизируется независимо, с совместной оптимизацией внутри группы. Этот подход улучшает компромисс между использованием кодовой книги и производительностью реконструкции. Кроме того, мы вводим метод переиспользования кодовой книги без обучения, позволяющий корректировать размер кодовой книги после обучения. В экспериментах по реконструкции изображений в различных условиях Group-VQ демонстрирует улучшенную производительность по метрикам реконструкции.

Исследовательский контекст и мотивация

Описание проблемы

Vector Quantization (VQ) — это метод отображения непрерывных признаков в дискретные токены, широко применяемый в VQ-VAE. Однако традиционное обучение VQ сталкивается с проблемой низкого использования кодовой книги, когда используются и обновляются только некоторые кодовые векторы, что приводит к "коллапсу кодовой книги" и ограничивает способность кодирования модели.

Ограничения существующих методов

1. Vanilla VQ: каждый кодовый вектор обновляется независимо, что легко приводит к коллапсу кодовой книги
2. Методы Joint VQ (такие как SimVQ, VQGAN-LC): совместная оптимизация всей кодовой книги посредством совместного использования параметров достигает 100% использования, но ограничивает способность обучения кодовой книги

Исследовательская мотивация

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

Основные вклады

1. Предложение метода Group-VQ: метод оптимизации кодовой книги на основе группировки, балансирующий использование и производительность реконструкции в моделях VQ
2. Обобщение метода Joint VQ: переосмысление Joint VQ с точки зрения совместного использования параметров и введение метода переиспользования кодовой книги после обучения
3. Корректировка кодовой книги без обучения: реализация гибкой корректировки размера кодовой книги после обучения без переобучения модели
4. Комплексная экспериментальная проверка: верификация эффективности Group-VQ и переиспользования кодовой книги на задачах реконструкции изображений

Детальное описание метода

Определение задачи

Для изображения $I \in \mathbb{R}^{H \times W \times 3}$ VQ-VAE сначала использует кодировщик для получения карты признаков $Z \in \mathbb{R}^{h \times w \times d}$, затем квантизатор заменяет каждый вектор признаков $z \in \mathbb{R}^d$ на ближайший кодовый вектор из кодовой книги $C = \{q_i | q_i \in \mathbb{R}^d, i = 0,1,...,n-1\}$:

$q = \arg\min_{q_i \in C} \|z - q_i\|, i = 0,1,...,n-1$

Архитектура модели

Дизайн Group-VQ

Group-VQ разделяет кодовую книгу $C$ на $k$ непересекающихся групп (подкодовых книг):

$C = \bigcup_{j=0}^{k-1} G_j, \quad G_j \cap G_{j'} = \emptyset \text{ если } j \neq j'$

Каждая группа $G_j$ обновляется независимо с совместной оптимизацией внутри группы. Для кодового вектора $q_{jt} \in G_j$ обновление градиента выглядит следующим образом:

$\nabla_{q_{jt}} L_{cmt} = \nabla_{q_{jt}} L_j$

Это гарантирует, что каждая группа получает градиенты только от кодовых векторов внутри неё.

Параметризация кодовой книги

Каждая группа $G_j$ параметризуется посредством совместного использования параметров:

$G_j = \hat{G}_j W_j + b_j$

где:

• $\hat{G}_j \in \mathbb{R}^{n_j \times r_j}$: ядро кодовой книги (фиксированное распределение выборки)
• $W_j \in \mathbb{R}^{r_j \times d}$: проектор (обучаемый)
• $b_j \in \mathbb{R}^d$: вектор смещения

Технические инновации

1. Анализ с единой точки зрения

• Vanilla VQ: $k = n$, каждый кодовый вектор — это одна группа
• Joint VQ: $k = 1$, вся кодовая книга — одна группа
• Group-VQ: $1 \leq k \leq n$, балансирование двух крайних случаев

2. Механизм переиспользования кодовой книги

Используя свойства генеративной кодовой книги, после обучения можно переиспользовать ядро кодовой книги:

$\tilde{q} = \hat{v} W_j, \quad \hat{v} \sim \mathcal{N}(0, I)$

Поддерживает два режима:

• Переиспользование: полная замена кодовой книги
• Самораспространение: добавление новых кодовых векторов к исходной кодовой книге

Экспериментальная установка

Наборы данных

• ImageNet-1k: основной набор данных
• MS-COCO: дополнительная проверка
• Разрешение входа: 128×128, коэффициент понижающей дискретизации f=8

Метрики оценки

• rFID (reconstruction FID): расстояние распределения между реконструированным и исходным изображением
• LPIPS(VGG16): перцептивная схожесть
• PSNR: пиковое отношение сигнал-шум
• SSIM: индекс структурного сходства

Методы сравнения

• VQGAN, ViT-VQGAN, VQGAN-FC
• FSQ, LFQ (методы с фиксированной кодовой книгой)
• VQGAN-LC, SimVQ (методы Joint VQ)

Детали реализации

• Скорость обучения: 1×10⁻⁴
• Оптимизатор: Adam (β₁=0.5, β₂=0.9)
• Размер пакета: 32/GPU
• Оборудование: NVIDIA A5000 GPU

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

Основные результаты

Сравнение производительности на ImageNet-1k (размер кодовой книги 65,536):

Group-VQ достигает лучшей производительности по всем метрикам, значительно превосходя методы-базовые линии.

Абляционные исследования

Влияние различного количества групп:

Эксперименты показывают, что 32-64 группы — оптимальный выбор, балансирующий использование кодовой книги и производительность реконструкции.

Эксперименты по переиспользованию кодовой книги

Результаты корректировки размера кодовой книги:

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

Визуальный анализ

Путём случайной проекции кодовых векторов в 2D-пространство обнаружено:

1. Различные группы изучают различные распределения признаков
2. Кодовые векторы внутри группы относительно похожи, различия между группами значительны
3. Статистические характеристики каждой группы (среднее, дисперсия, частота использования) существенно отличаются

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

Классификация методов улучшения VQ

1. Улучшения Straight-Through Estimator: оптимизация распространения градиентов
2. Многоиндексное квантование: RQ-VAE, Product Quantization и др.
3. Улучшения кодовой книги: основной фокус данной работы

Методы Joint VQ

• VQGAN-LC: использование предварительно обученных признаков для инициализации + слой проекции
• SimVQ: случайная инициализация + переиспользование матрицы
• LFQ/FSQ: фиксированная кодовая книга для предотвращения коллапса

В данной работе эти методы объединены как "Joint VQ, реализованный посредством совместного использования параметров", и на этой основе предложена стратегия групповой оптимизации.

Заключение и обсуждение

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

1. Компромисс между использованием кодовой книги и качеством реконструкции: 100% использование не обязательно приводит к лучшему результату реконструкции
2. Групповая оптимизация — эффективная стратегия балансировки: Group-VQ обеспечивает гибкий контроль путём корректировки количества групп
3. Переиспользование кодовой книги имеет практическую ценность: после обучения можно гибко корректировать размер кодовой книги

Ограничения

1. Отсутствие проверки на генеративных задачах: тестирование только на задачах реконструкции, отсутствует проверка на генеративных моделях
2. Требуется настройка количества групп: оптимальное количество групп зависит от конкретной задачи и набора данных
3. Вычислительная сложность: оптимизация нескольких групп может увеличить время обучения

Направления будущих исследований

1. Проверка эффективности Group-VQ на генеративных моделях (таких как авторегрессивные модели)
2. Исследование стратегий адаптивного выбора количества групп
3. Изучение комбинирования Group-VQ с другими методами улучшения VQ

Глубокая оценка

Преимущества

1. Ясный теоретический вклад: объединение существующих методов VQ с точки зрения групповой оптимизации, предоставление нового аналитического подхода
2. Простой и эффективный метод: дизайн Group-VQ интуитивен, легко реализуется и понимается
3. Полные эксперименты: комплексная проверка на нескольких наборах данных и архитектурах, детальные абляционные исследования
4. Высокая практическая ценность: метод переиспользования кодовой книги решает потребность в гибкости при практическом применении

Недостатки

1. Недостаточно глубокий теоретический анализ: отсутствует теоретическое объяснение того, почему групповая оптимизация более эффективна
2. Ограниченная область применения: основное внимание на реконструкции изображений, эффективность на других модальностях и задачах неизвестна
3. Отсутствует анализ вычислительных затрат: не проведён детальный анализ вычислительных расходов оптимизации нескольких групп

Влияние

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

Сценарии применения

1. Кодирование изображений/видео: задачи сжатия, требующие высокого качества реконструкции
2. Мультимодальное обучение: как универсальный компонент векторного квантования
3. Генеративные модели: как токенизатор для предоставления дискретного представления генеративным моделям

Библиография

Данная работа основана на следующих важных исследованиях:

1. Van Den Oord et al. (2017) — оригинальная статья VQ-VAE
2. Zhu et al. (2024b) — метод SimVQ
3. Yu et al. (2023) — метод LFQ
4. Mentzer et al. (2023) — метод FSQ

Резюме: Это статья с важным вкладом в область VQ. Метод Group-VQ простой и эффективный, предоставляет новый подход к оптимизации кодовой книги. Метод переиспользования кодовой книги имеет очень сильную практическую ценность. Хотя есть место для улучшения в теоретическом анализе и области применения, в целом это высокое качество исследовательской работы.