Эффективное генеративное моделирование с использованием токенов на основе остаточной векторной квантизации

В последние годы мы наблюдаем значительный прогресс в области глубоких генеративных моделей, которые успешно генерируют высококачественные и реалистичные данные в различных областях, таких как языковое моделирование, генерация изображений и синтез аудио. Однако, несмотря на достижения, существует множество вызовов, связанных с обеспечением высокой достоверности генерации, особенно когда речь идет о балансе между качеством генерации и вычислительной эффективностью. В этой статье рассматривается метод, называемый Residual Vector Quantization (RVQ), который улучшает качество реконструкции данных без увеличения длины последовательности.

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

В данной работе мы представляем ResGen — эффективную генеративную модель на основе RVQ, которая достигает высокой достоверности генерации без компромиссов в скорости выборки. Основная идея заключается в прямом предсказании векторных встраиваний коллективных токенов, а не предсказании каждого токена по отдельности. Это позволяет нам декомпозировать сложность выборки как по длине последовательности, так и по глубине, что приводит к модели, которая эффективно генерирует высококачественные образцы.

Основные концепции

Остаточная векторная квантизация (RVQ)

RVQ представляет собой метод, который улучшает VQ-VAEs, применяя векторную квантизацию к остаткам предыдущих квантизаций. В этом процессе выходные данные энкодера в VQ-VAE на позиции i обозначаются как ( h_{i,0} ). Остаточный векторный квантизатор отображает это значение в последовательность квантизированных токенов ( x \in \mathbb{N}^{L \times D} ), где ( D ) — общая глубина процесса RVQ.

Каждый токен ( x_{i,j} ) выбирается на основе ближайшего вектора из кодовой книги, что позволяет RVQ эффективно захватывать наиболее значимые характеристики в нижних слоях квантизации, в то время как более тонкие детали постепенно захватываются в более высоких слоях.

Генеративное моделирование с маскированными токенами

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

Метод ResGen

ResGen строится на основе вышеописанных концепций, где мы вводим стратегию маскирования токенов, специально адаптированную для токенов RVQ. Мы моделируем предсказание маскированных токенов, предсказывая сумму векторных встраиваний, чтобы декомпозировать итерации генерации от длины и глубины последовательностей токенов.

Процесс маскирования и предсказания

Процесс маскирования токенов начинается с высших уровней квантизации, что позволяет эффективно использовать иерархическую природу RVQ. Мы применяем бинарную маску, где каждый элемент указывает, маскирован ли токен или нет. Общее количество токенов для маскирования определяется по заранее заданной маскировочной функции.

Формулировка в рамках вероятностной модели

Для обеспечения теоретической основы нашего метода, мы рассматриваем процесс предсказания маскированных токенов как шаги прямого диффузионного процесса. Это позволяет нам интерпретировать маскированный процесс как шаги диффузии, постепенно преобразуя полностью немаскированную последовательность токенов в полностью маскированную.

Экспериментальные результаты

Мы провели эксперименты по генерации изображений и синтезу речи, чтобы продемонстрировать эффективность нашего подхода. Для задач в области зрения мы сосредоточились на условной генерации изображений на наборе данных ImageNet с разрешением 256×256. В области аудио мы оценили нашу модель с использованием задач продолжения речи и кросс-предложения.

Результаты показывают, что ResGen превосходит авторегрессионные модели в обеих задачах, обеспечивая высокое качество генерации при меньших затратах времени на выборку. Кроме того, по мере увеличения глубины RVQ, наша модель демонстрирует улучшение качества генерации или скорость выборки по сравнению с аналогичными моделями.

Заключение

В данной работе мы представили ResGen — эффективную генеративную модель на основе RVQ, которая генерирует высококачественные образцы, сохраняя при этом быструю скорость выборки. Наша методология включает в себя предсказание векторных встраиваний коллективных токенов, что позволяет эффективно справляться с типичными компромиссами между глубиной токенов и скоростью выборки в векторно-квантизированных генеративных моделях. Эксперименты подтвердили высокую производительность ResGen, которая сопоставима или превосходит авторегрессионные модели по качеству и скорости выборки.

Будущие направления

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