👀 Robust Training of Vector Quantized Bottleneck Models

Сейчас обсудим работу, дающую полезную интуицию про обучение VQ-VAE. Предыдущий пост был про основные проблемы, а этот будет про их решения. Больше всего внимания было уделено "коллапсу кодбука" [1], [2], [3], [4], [5], [6], [7], явление при котором используется меньшая часть элементов кодбука. Еще в первой статье про VQ-VAE [1] была предложена стратегия обновления кодбука с использованием EMA. Следующая работа [2] изучает стратегию с EMA и попутно предлагает новое решение.

Авторы утверждают, что норма элементов кодбука должна быть меньше нормы feature maps из энкодера. Чтобы сохранять масштаб (норму) элементов кодбука, будет использоваться батч нормализация.

Другая важная мысль будет о том, что K-Means это алгоритм рассчитанный на фиксированные данные. То есть, при обучении кодбука используются feature maps, которые меняются от семпла к семплу. Обновление кластеров может быть слишком быстрым, чтобы адаптировать все элементы кодбука. Особенно при том, что градиенты влияют только на последние полученные элементы кодбука.

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

Далее операция VQ включается в процесс обучения и через K-Means++ создается кодбук. Кластеризация будет происходить по семплам собранным через reservoir sampling [8]. Это будет гарантировать, что в начале все элементы кодбука будут использоваться. Затем, чтобы адаптироваться к изменениям в распределении выходов энкодера, кодовая книга будет периодически обновляться по той же схеме. Ну а за метриками отправляйтесь в статью. Методы не самый актуальный, но идеи были хорошие. Продолжим обсуждение других работ в следующих постах.

#papers #vqvae #audio #images