Cuantización Escalar Finita: VQ-VAE Simplificado

Finite Scalar Quantization: VQ-VAE Made Simple

September 27, 2023

Autores: Fabian Mentzer, David Minnen, Eirikur Agustsson, Michael Tschannen

cs.AI

Resumen

Proponemos reemplazar la cuantización vectorial (VQ) en la representación latente de los VQ-VAE con un esquema simple denominado cuantización escalar finita (FSQ), donde proyectamos la representación del VAE en unas pocas dimensiones (típicamente menos de 10). Cada dimensión se cuantiza a un pequeño conjunto de valores fijos, lo que da lugar a un (implícito) libro de códigos dado por el producto de estos conjuntos. Al elegir adecuadamente el número de dimensiones y los valores que cada dimensión puede tomar, obtenemos el mismo tamaño de libro de códigos que en VQ. Sobre tales representaciones discretas, podemos entrenar los mismos modelos que se han entrenado con representaciones de VQ-VAE. Por ejemplo, modelos autoregresivos y transformadores enmascarados para generación de imágenes, generación multimodal y tareas de visión por computadora de predicción densa. Concretamente, empleamos FSQ con MaskGIT para generación de imágenes, y con UViM para estimación de profundidad, colorización y segmentación panóptica. A pesar del diseño mucho más simple de FSQ, obtenemos un rendimiento competitivo en todas estas tareas. Destacamos que FSQ no sufre de colapso del libro de códigos y no necesita la maquinaria compleja empleada en VQ (pérdidas de compromiso, resiembra del libro de códigos, división de códigos, penalizaciones de entropía, etc.) para aprender representaciones discretas expresivas.

English

We propose to replace vector quantization (VQ) in the latent representation of VQ-VAEs with a simple scheme termed finite scalar quantization (FSQ), where we project the VAE representation down to a few dimensions (typically less than 10). Each dimension is quantized to a small set of fixed values, leading to an (implicit) codebook given by the product of these sets. By appropriately choosing the number of dimensions and values each dimension can take, we obtain the same codebook size as in VQ. On top of such discrete representations, we can train the same models that have been trained on VQ-VAE representations. For example, autoregressive and masked transformer models for image generation, multimodal generation, and dense prediction computer vision tasks. Concretely, we employ FSQ with MaskGIT for image generation, and with UViM for depth estimation, colorization, and panoptic segmentation. Despite the much simpler design of FSQ, we obtain competitive performance in all these tasks. We emphasize that FSQ does not suffer from codebook collapse and does not need the complex machinery employed in VQ (commitment losses, codebook reseeding, code splitting, entropy penalties, etc.) to learn expressive discrete representations.