👀 Scaling the Codebook Size of VQGAN to 100,000 with a Utilization Rate of 99%
[статья] | [код]

Давайте сегодня закончим тему колапса кодбука и последнее слово будет за VQGAN-LC (Large Codebook), свежая работа на NIPS 2024. Вы видите??? 100 тыщ кодбук!

Идея мне очень понравилась. Помните я говорил, что колапс кодбука происходит потому, что мы обновляем кодбук относительно последних фичермапов. Некоторые элементы кодбука могут остаться без изменений и не будут использоваться. Данное решение примечательно тем, что каждый раз обновляется весь кодбук. Теперь мы стараемся приблизить фичермапы и кодбук не по семплам, а приближаем распределение кодбука к фичермапам.

Кодбук будет фиксированным, инициализирован с помощью CLIP энкодера с ViT бекбоном. Далее делаем проекцию кодбука через линейный слой и считаем L2 между проекцией и элементами фичермапов, получаем квантованые фичермапы. Затем считаем стандартный VQ лосс, но кодбук остается фиксированным, а линейный слой для проекции кодбука и наш энкодер обучаются.

Получается, что теперь мы учитываем все элементы кодбука на каждом степе и учимся распределять фичермапы максимально эффективно вплоть до кодбука размером 100к! И до миллиона можно разогнать, если декодер будет побольше.

Но если посмотреть на результаты, то далеко не всегда размер кодбука значительно улучшает качество. Больше не всегда хорошо, здесь, конечно, об этом не написано. Ну, все же чаще масштабирование в большую сторону помогает :)

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

#papers #vqvae #images