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.