VQ-Diffusion

'VQ-Diffusion'

La Difusión Vector Cuantizada (VQ-Diffusion) es un modelo de difusión latente condicional desarrollado por la Universidad de Ciencia y Tecnología de China y Microsoft. A diferencia de la mayoría de los modelos de difusión estudiados comúnmente, los procesos de ruido y desruido de VQ-Diffusion operan en un espacio latente cuantizado, es decir, el espacio latente está compuesto por un conjunto discreto de vectores. Los modelos de difusión discreta se han explorado menos que sus contrapartes continuas y ofrecen un punto interesante de comparación con los modelos autorregresivos (AR).

• Tarjeta de modelo de Hugging Face
• Espacios de Hugging Face
• Implementación original
• Artículo

Demo

🧨 Diffusers te permite ejecutar VQ-Diffusion con solo unas pocas líneas de código.

Instalar dependencias

pip install 'diffusers[torch]' transformers ftfy

Cargar el pipeline

• Pronóstico de series de tiempo probabilísticas con 🤗 Transformers
• Usando la Difusión Estable con Core ML en Apple Silicon
• De GPT2 a Stable Diffusion Hugging Face llega a la comunidad de Elixir

from diffusers import VQDiffusionPipeline

pipe = VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq")

Si quieres utilizar pesos FP16

from diffusers import VQDiffusionPipeline
import torch

pipe = VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq", torch_dtype=torch.float16, revision="fp16")

Mover a la GPU

pipe.to("cuda")

¡Ejecutar el pipeline!

prompt = "Un osito de peluche jugando en la piscina."

image = pipe(prompt).images[0]

Arquitectura

VQ-VAE

Las imágenes se codifican en un conjunto de “tokens” o vectores de incrustación discretos utilizando un codificador VQ-VAE. Para hacerlo, las imágenes se dividen en parches y luego cada parche se reemplaza por la entrada más cercana de un libro de códigos con un vocabulario de tamaño fijo. Esto reduce la dimensionalidad del espacio de píxeles de entrada. VQ-Diffusion utiliza la variante VQGAN de Taming Transformers . Esta publicación de blog es un buen recurso para comprender mejor los VQ-VAEs.

VQ-Diffusion utiliza un VQ-VAE preentrenado que se congeló durante el proceso de entrenamiento de difusión.

Proceso hacia adelante

En el proceso de difusión hacia adelante, cada token latente puede mantenerse igual, ser muestreado nuevamente a un vector latente diferente (cada uno con igual probabilidad) o ser enmascarado. Una vez que un token latente está enmascarado, se mantendrá enmascarado. α t \alpha_t α t ​ , β t \beta_t β t ​ y γ t \gamma_t γ t ​ son hiperparámetros que controlan el proceso de difusión hacia adelante desde el paso t − 1 t-1 t − 1 hasta el paso t t t . γ t \gamma_t γ t ​ es la probabilidad de que un token sin enmascarar se vuelva enmascarado. α t + β t \alpha_t + \beta_t α t ​ + β t ​ es la probabilidad de que un token sin enmascarar se mantenga igual. El token puede transicionar a cualquier vector latente no enmascarado individual con una probabilidad de β t \beta_t β t ​ . En otras palabras, α t + K β t + γ t = 1 \alpha_t + K \beta_t + \gamma_t = 1 α t ​ + K β t ​ + γ t ​ = 1 donde K K K es el número de vectores latentes no enmascarados. Consulta la sección 4.1 del artículo para obtener más detalles.

Aproximando el proceso inverso

Un transformador codificador-decodificador aproxima las clases de los latentes sin ruido, x 0 x_0 x 0 ​ , condicionados al estímulo, y y y . El codificador es un codificador de texto CLIP con pesos congelados. El transformador del decodificador proporciona atención global sin enmascaramiento a todos los píxeles latentes y produce las probabilidades logarítmicas de la distribución categórica sobre las incrustaciones de vectores. El transformador del decodificador predice toda la distribución de latentes sin ruido en un solo paso hacia adelante, proporcionando autoatención global sobre x t x_t x t ​ . Formular el problema como una secuencia condicional a secuencia de valores discretos proporciona cierta intuición sobre por qué el transformador codificador-decodificador es adecuado.

La sección de modelos AR proporciona contexto adicional sobre la arquitectura de VQ-Diffusion en comparación con los modelos basados en transformadores AR.

Domar a los Transformadores ofrece una buena discusión sobre cómo convertir píxeles en tokens discretos en un espacio latente comprimido para que los transformadores sean computacionalmente factibles para datos de imágenes.

VQ-Difusión en Contexto

Modelos de Difusión

Los modelos de difusión contemporáneos son en su mayoría continuos. En el proceso hacia adelante, los modelos de difusión continuos añaden ruido gaussiano de forma iterativa. El proceso inverso se aproxima a través de p θ ( x t − 1 ∣ x t ) = N ( x t − 1 ; μ θ ( x t , t ) , Σ θ ( x t , t ) ) p_{\theta}(x_{t-1} | x_t) = N(x_{t-1}; \mu_{\theta}(x_t, t), \Sigma_{\theta}(x_t, t)) p θ ​ ( x t − 1 ​ ∣ x t ​ ) = N ( x t − 1 ​ ; μ θ ​ ( x t ​ , t ) , Σ θ ​ ( x t ​ , t ) ) . En el caso más simple de DDPM , la matriz de covarianza es fija, se entrena un U-Net para predecir el ruido en x t x_t x t ​ , y x t − 1 x_{t-1} x t − 1 ​ se obtiene a partir del ruido.

El proceso inverso aproximado es estructuralmente similar al proceso inverso discreto. Sin embargo, en el caso discreto, no hay un análogo claro para predecir el ruido en x t x_t x t ​ , y predecir directamente la distribución para x 0 x_0 x 0 ​ es un objetivo más claro.

Hay una menor cantidad de literatura que abarca modelos de difusión discreta en comparación con modelos de difusión continua. Aprendizaje profundo no supervisado utilizando termodinámica de no equilibrio presenta un modelo de difusión sobre una distribución binomial. Argmax Flows y Multinomial Diffusion extiende la difusión discreta a distribuciones multinomiales y entrena un transformador para predecir la distribución sin ruido para una tarea de modelado de lenguaje. Modelos de Difusión con Desruido Estructurado en Espacios de Estados Discretos generaliza la difusión multinomial con procesos de desruido alternativos: uniforme, absorbente, gaussiano discretizado y distancia de incrustación de token. Los procesos de desruido alternativos también son posibles en modelos de difusión continua, pero como se señala en el artículo, solo el ruido gaussiano aditivo ha recibido una atención significativa.

Modelos Autoregresivos

Quizás sea más interesante comparar VQ-Difusión con los modelos AR, ya que con mayor frecuencia presentan transformadores que hacen predicciones sobre distribuciones discretas. Si bien los transformadores han demostrado éxito en el modelado AR, aún sufren disminuciones lineales en la velocidad de inferencia para una mayor resolución de imágenes, acumulación de errores y sesgo direccional. VQ-Difusión mejora en estos tres puntos problemáticos.

Los modelos generativos de imágenes AR se caracterizan por factorizar la probabilidad de la imagen de manera que cada píxel esté condicionado a los píxeles anteriores en un orden de escaneo de ráster (de izquierda a derecha, de arriba a abajo), es decir, p ( x ) = ∏ i p ( x i ∣ x i − 1 , x i − 2 , . . . x 2 , x 1 ) p(x) = \prod_i p(x_i | x_{i-1}, x_{i-2}, … x_{2}, x_{1}) p ( x ) = ∏ i ​ p ( x i ​ ∣ x i − 1 ​ , x i − 2 ​ , . . . x 2 ​ , x 1 ​ ) . Como resultado, los modelos pueden entrenarse maximizando directamente la log-verosimilitud. Además, los modelos AR que operan en valores de píxeles reales (no latentes) predicen valores de canal a partir de una distribución multinomial discreta, es decir, primero se muestrea el valor del canal rojo de un softmax de 256 vías, y luego la predicción del canal verde se condiciona al valor del canal rojo.

Los modelos generativos de imágenes AR han evolucionado arquitectónicamente con mucho trabajo para hacer que los transformadores sean computacionalmente factibles. Antes de los modelos basados en transformadores, PixelRNN, PixelCNN y PixelCNN++ eran el estado del arte.

Transformador de Imágenes ofrece una buena discusión sobre los modelos no basados en transformadores y la transición a modelos basados en transformadores (consultar el artículo para ver las citas omitidas).

Entrenar redes neuronales recurrentes para predecir secuencialmente cada píxel, incluso en una imagen pequeña, es computacionalmente muy desafiante. Por lo tanto, los modelos paralelizables que utilizan redes neuronales convolucionales, como el PixelCNN, han recibido recientemente mucha más atención y ahora han superado al PixelRNN en calidad.

Una desventaja de las CNN en comparación con las RNN es su campo receptivo típicamente bastante limitado. Esto puede afectar adversamente su capacidad para modelar fenómenos a larga distancia comunes en las imágenes, como la simetría y la oclusión, especialmente con un número reducido de capas. Se ha demostrado que aumentar el campo receptivo mejora significativamente la calidad (Salimans et al.). Sin embargo, hacerlo conlleva un costo significativo en términos de número de parámetros y, en consecuencia, de rendimiento computacional, lo que puede dificultar el entrenamiento de dichos modelos.

… la autoatención puede lograr un mejor equilibrio en el compromiso entre el campo receptivo virtualmente ilimitado del PixelRNN, que necesariamente es secuencial, y el campo receptivo limitado del PixelCNN mucho más paralelizable y sus diversas extensiones.

Image Transformer utiliza transformadores restringiendo la autoatención a vecindarios locales de píxeles.

Taming Transformers y DALL-E 1 combinan convoluciones y transformadores. Ambos entrenan un VQ-VAE para aprender un espacio latente discreto, y luego se entrena un transformador en el espacio latente comprimido. El contexto del transformador es global pero enmascarado, ya que se proporciona atención sobre todos los píxeles latentes predichos previamente, pero el modelo sigue siendo AR, por lo que no se puede proporcionar atención sobre los píxeles aún no predichos.

ImageBART combina convoluciones, transformadores y procesos de difusión. Aprende un espacio latente discreto que se comprime aún más con un proceso de difusión multinomial corto. Luego se entrenan transformadores codificador-decodificador separados para revertir cada paso en el proceso de difusión. El transformador codificador proporciona contexto global en x t x_t x t ​ mientras que el transformador decodificador predice autoregresivamente los píxeles latentes en x t − 1 x_{t-1} x t − 1 ​ . Como resultado, cada píxel recibe atención cruzada global en la imagen más ruidosa. Se utilizan entre 2 y 5 pasos de difusión con más pasos para conjuntos de datos más complejos.

A pesar de haber logrado avances tremendos, los modelos AR todavía sufren de disminuciones lineales en la velocidad de inferencia para una mayor resolución de imagen, acumulación de errores y sesgo direccional. Para modelos de transformadores AR de tamaño equivalente, el big-O de la inferencia de VQ-Diffusion es mejor siempre y cuando el número de pasos de difusión sea menor que el número de píxeles latentes. Para el conjunto de datos ITHQ, la resolución latente es de 32×32 y el modelo se entrena hasta 100 pasos de difusión para una mejora de aproximadamente 10 veces en el big-O. En la práctica, VQ-Diffusion “puede ser 15 veces más rápido que los métodos AR mientras logra una mejor calidad de imagen” (consulte el artículo para obtener más detalles). Además, VQ-Diffusion no requiere teacher-forcing y en su lugar aprende a corregir tokens predichos incorrectamente. Durante el entrenamiento, las imágenes con ruido se enmascaran y los píxeles latentes se reemplazan por tokens aleatorios. VQ-Diffusion también puede proporcionar contexto global en x t x_t x t ​ mientras predice x t − 1 x_{t-1} x t − 1 ​ .

Pasos adicionales con VQ-Diffusion y 🧨 Diffusers

Hasta ahora, solo hemos adaptado el modelo VQ-Diffusion entrenado en el conjunto de datos ITHQ. También se han lanzado modelos VQ-Diffusion entrenados en CUB-200, Oxford-102, MSCOCO, Conceptual Captions, LAION-400M e ImageNet.

VQ-Diffusion también admite una estrategia de inferencia más rápida. La reproyección de la red se basa en la posterior del proceso de difusión condicionado a la imagen sin ruido, que es tratable. Una fórmula similar se aplica al usar un paso de tiempo, Δ t \Delta t Δ t , que salta un número de pasos de difusión inversa, p θ ( x t − Δ t ∣ x t , y ) = ∑ x ~ 0 = 1 K q ( x t − Δ t ∣ x t , x ~ 0 ) p θ ( x ~ 0 ∣ x t , y ) p_\theta (x_{t – \Delta t } | x_t, y) = \sum_{\tilde{x}_0=1}^{K}{q(x_{t – \Delta t} | x_t, \tilde{x}_0)} p_\theta(\tilde{x}_0 | x_t, y) p θ ​ ( x t − Δ t ​ ∣ x t ​ , y ) = ∑ x ~ 0 ​ = 1 K ​ q ( x t − Δ t ​ ∣ x t ​ , x ~ 0 ​ ) p θ ​ ( x ~ 0 ​ ∣ x t ​ , y ) .

Improved Vector Quantized Diffusion Models mejora la calidad de las muestras de VQ-Diffusion con una guía discreta sin clasificador y una estrategia de inferencia alternativa para abordar el “problema de la distribución conjunta” — consulte la sección 3.2 para más detalles. La guía discreta sin clasificador se fusiona en los difusores, pero aún no se ha agregado la estrategia de inferencia alternativa.

¡Se agradecen las contribuciones!

We will continue to update Zepes; if you have any questions or suggestions, please contact us!