VQ-VAE: когда нужны дискретные представления и как их готовить

VQ-VAE — популярная нейросетевая архитектура, которая активно применяется в моделях, работающих с аудио, изображениями, видео, в мультимодальных моделях, а также в 3D-моделировании, трекинге движений и естественных науках (биологии, химии). На основе VQ-VAE создаются модели для сжатия и восстановления данных, которые сейчас используются во многих стриминговых платформах. Также VQ-VAE применяется для задач порождения или извлечения информации; например, на нём работает моё любимое audio2midi, которое я делаю, или panoptic segmentation.

В общем, идея полезная, но у VQ-VAE есть конкретные сценарии применения и проблемы, о которых речь пойдёт далее:
— он вычисляет N центроидов кластеров по картам признаков, полученным из кодировщика; кластеризация приводит к потере информации;
— инициализация центроидов кластеров значительно влияет на метрики для любой задачи; всем известный k-means даёт здесь полезную последовательность идей;
— через операцию VQ не протекают градиенты, поэтому мы вынуждены использовать аппроксимации, чаще всего STE;
— поскольку градиент останавливается после VQ, он не доходит до словаря и кодировщика; нужны вспомогательные функции потерь. Одна из них будет приближать выходы кодировщика к элементам словаря, вторая — приближать элементы словаря к выходам кодировщика; а это приводит к сложностям в оптимизации.

Поэтому для некоторых моделей в CV используют просто патчи, как в ViT. Это тоже интересный момент, и мы ещё вернемся к нему, а пока идём дальше.

— Часто вспомогательные лоссы считают L2-расстояние между конкретными семплами, и возникает коллапс словаря, когда модель использует только небольшую часть токенов, игнорируя остальные. Это снижает эффективность модели и происходит потому, что мы приближаем друг к другу лишь конкретные семплы, и ближе всех к выходам кодировщика будут последние, а большинство предыдущих окажутся дальше последних и никогда не будут выбраны. Поэтому создавать большие словари (10К+) не имеет смысла.

— При обучении с нуля подбор числа элементов словаря и их размерности — это трудоёмкий процесс: каждый раз нужно инициализировать словарь заново и проверять разные центроиды кластеров.

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

— Часто для задач порождения или извлечения информации используют предобученные кодеки, и никто не хочет их менять под данные, модель и задачу, т.к. это может испортить кодек и потребует долгого подбора параметров. Но страдать придётся в любом случае: нужно подбирать и оптимизировать декодер на основе токенов из выбранного кодека.

— Регионы, по которым вычисляются латентные представления, имеют разную информационную плотность; в частности, мы хотим разделять регионы по смыслу, например токенизировать фон на изображении отдельно от основных объектов.

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