Finite Scalar Quantization: VQ-VAE Made Simple
[Статья][Код скорее есть, чем нет]

Очередная статья про квантизацию, но в другом контексте 🤪.
Вариационные автокодировщики с дискретными латентными кодами показали себя в качесте неплохих генеративных моделей (VQ-VAE, VQ-VAE-2). Кроме того, их способность эффективно представлять высокоразмерные данные лежит в основе латентных диффузионных моделей.

Однако, обучение VQ-VAE - довольно нетривиальная задача, требующая тонкой настройки. В отсутствие специальных манипуляций активации отображаются в малое число кодовых векторов, в то время как основная масса кодов сидит и грустит 😞. Приходится накладывать регуляризацию, чтобы равномерно заметать кодовую книгу. Коды обновляются через EMA. Поиск ближайшего соседа также требует определенных вычислительных затрат.

В этой работе предложили использовать фиксированную гиперкубическую решетку для квантования представлений.

Метод

Входные активации отображаются в некоторый гиперкуб с заданным числом узлов L вдоль каждой размерности. Чтобы ограничить диапазон принимаемых значений, ко входу применяют tanh:

f(z) = [L/2] tanh(z) [] - округление к ближайшему целому

Кодовые векторы - всевозможные узлы в этой решетке, т.е размер кодовой книги L^d, d - размерность латентного пространства. Причем размерность пространства небольшая - не более 10 во всех экспериментах.

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

Эксперименты

Метод валидируют на архитектурах MaskGiT (генерация изображений) и UViM (оценка глубины). На малом числе кодовых векторов стандартная формулировка работает немного лучше, с увеличением числа кодов VQ-VAE общего вида начинает страдать в качестве, в то время как у FSQ метрики монотонно уменьшаются.

По FID оптимальное качество FSQ на одном уровне с VQ. FSQ задействует все кодовые векторы, в то время как у базового VQ есть неиспользуемые коды.

На depth estimation метрики также на одном уровне с VQ-VAE общего вида.

Вывод

Занятно, что в области VQ-VAE пошли в противоположную сторону от текущих трендов квантизации LLM - от более общей векторной квантизации к скалярной. Интересно, насколько полученные коды хороши для обучения латентных диффузионных моделей?