👀 SQ-VAE: Variational Bayes on Discrete Representation with Self-annealed Stochastic Quantization
[Статья] | [Код]

Напомню, насколько важен VQ-VAE и продолжу делать заметки про коллапс кодбука. Явление, при котором модель использует лишь небольшую часть доступных кодов, что ограничивает её способность эффективно обрабатывать данные. Сегодня обсудим SQ-VAE, частично решающий колапс кодбука. SQ-VAE обучается генерации без вспомогательных лоссов для кодбука и энкодера, то есть градиент без проблем проходит через операцию векторного квантования.

Авторы решили вернуться к истокам, точнее к VAE, и сделали процесс квантования стохастическим. Фичермапы семплируются из непрерывного распределения с обучаемыми параметрами, например из Гауссова (рис. 1). Теперь, как в VAE, мы оптимизируем ELBO и можем применить reparametrization trick, так операция семплирования не остановит градиент.

При этом под конец обучения стохастичность приближается к нулю и мы возвращаемся к детерминированному векторному квантованию. Это явление авторы назвали self-annealing (рис. 2, эксперимент на mnist). Довольно удобный результат, на инференсе не будет семплирования и мы вернемся к классическому VQ-VAE. Такой сценарий наблюдается если распределение данных непрерывное.

Что если распределение данных дискретное? Интуитивный способ адаптации SQ-VAE для категориального распределения данных — смоделировать выходных данные декодера в виде категориального распределения. Но при таком сценарии self-annealing не наблюдается. Чтобы исправить это авторы используют von Mises-Fisher distribution для стохастического векторного квантования. Истинная причина проблемы лежит в исчезновении дисперсии в лоссе после замены непрерывного распределения на дискретное (см. статью в 3.4).

Метрики по картиночным и аудио задачам на рис. 3 и работа с 2022, поэтому посмотрите цитирования.

#papers #vqvae #images #audio