[Resumen] Transformadores Adiestrados para la Síntesis de las Imágenes de Alta Resolución

La eficiencia de las metodologías convolucionales con la expresividad de los transformadores.
Created on March 5|Last edited on February 17
Comment
Traducción del artículo aquí.
﻿
﻿
 Este reporte presenta un resumen del paper “Transformadores Adiestrados para la Síntesis de las Imágenes de Alta Resolución”. Este paper apunta a los desafíos fundamentales de usar la expresividad de los transformadores para la síntesis de las imágenes de alta resolución. La metodología propuesta representa a las imágenes como a una composición de componentes de imágenes continuamente enriquecidos y, a su vez, utiliza a los transformadores para modelar eficientemente su composición con imágenes de alta resolución.
Aquí hay un hilo interesante de Twitter, hecho por Sander Dielemann, resumiendo a este paper.
﻿Paper |﻿ Página del Projecto| GitHub﻿Verifica mi Notebook de Colab para generar imágenes impresionantes como esta:
﻿
Run set1
﻿
IntroducciónLas Redes Neuronales Convolucionales (CNNs) son los modelos a utilizar para las tareas de visión. Esto se debe a que las CNNs exhiben un fuerte sesgo de localización (debido al uso de los kernels), y un sesgo hacia la invarianza espacial a través del uso de los pesos compartidos (un kernel barre la imagen completa).
En comparación con las CNNs, los transformadores no contienen un sesgo inductivo que priorice las interacciones locales, lo que les permite aprender relaciones complejas entre las entradas, en otras palabras haciéndolos expresivos. Sin embargo, esto es computacionalmente caro para secuencias a largo término. La expresividad incrementada de los transformadores viene con un incremento cuadrático de los costos computacionales.
Esser et al. (2021) demuestran la síntesis de las imágenes de alta resolución al combinar la efectividad de las CNNs con la expresividad de los transformadores.
Dos Palabras sobre los Modelos GenerativosDado  x≈pdata(x)x \approx p_{data}(x)x≈pdata​(x)﻿,  donde  pdata(x)p_{data}(x)pdata​(x)﻿ es la distribución verdadera que describe al conjunto de datos et X,X,X,﻿, el conjunto de datos consiste de muestras finitas a partir de esta distribución,
﻿X=x∣x≈pdata(x)X = {x|x \approx p_{data}(x)}X=x∣x≈pdata​(x)﻿﻿
La tarea generativa es encontrar un modelo tal que hat pmodel(x;θ)≈pdata(x)p_{model}(x; θ) \approx p_{data}(x)pmodel​(x;θ)≈pdata​(x)﻿. Aquí  θθθ﻿  son los parámetros del modelo.
Figura 1: Taxonomía del Modelo Generativo (Fuente)
Podemos distinguir dos tipos principales de modelos generativos:
Modelos basados en probabilidades (explícitos): Estos modelos proveen una especificación paramétrica “explícita” de la distribución de los datos y tienen una función de probabilidad manejable. Por ejemplo, los Autocodificadores Variacionales (VAEs) y los modelos autorregresivos.
Modelos implícitos: estos modelos no especifican la distribución de los datos en sí, sino que en su lugar definen un proceso estocástico que, después del entrenamiento, apunta a extraer muestras de la distribución de los datos subyacentes. Por ejemplo, Modelos Antagónicos Generativos (GANs).
Aquí hay algunas ideas claves para entender mejor a este paper:
 Las GANs, siendo implícitas, son difíciles de evaluar y, normalmente, fallan en la cobertura de todos los “modos” de los datos y, de esta forma, son más susceptibles al colapso modal.
Las GANs de mayor escala ahora pueden generar imágenes de alta calidad y de alta resolución. Sin embargo, es bien conocido que las muestras de estos modelos no capturan por completo la diversidad de la distribución verdadera. (Fuente)
Los métodos basados en probabilidades optimizan la probabilidad logarítmica negativa (NLL) de los datos del entrenamiento, lo que permite una comparación de los modelos más fácil y una mejor generalización de los datos no vistos. Sin embargo, maximizar la probabilidad en el espacio de los píxeles es complejo y computacionalmente caro.
 En un modelo autorregresivo, asumimos que un ejemplo x ∈ X puede ser representado como secuencias xt xix_ixi​﻿. La distribución es factorizada en un producto de condicionales, usando la regla de la cadena de la probabilidad dada como: : pdata(x)=q(x1)q(x2∣x1)..q(xd∣xd−1,...,x1).p_{data}(x) = q(x_1)q(x_2|x_1)..q(x_d|x_{d-1},..., x_1).pdata​(x)=q(x1​)q(x2​∣x1​)..q(xd​∣xd−1​,...,x1​).﻿ El modelo predice al próximo xi, basándose en el x<ix_{<i}x<i​﻿ previo . La generación de imágenes ha sido proyectada exitosamente como una generación de secuencias autorregreavisas o como un problema transformacional. (Fuente)
 Las GANs condicionales (cGANs) son una modificación simple, y aún efectiva, de una GAN regular, que permiten la generación de imágenes basándose en condiciones tales como la etiqueta de la clase u otra imagen (imagen parcial, mapa de segmentación, etc.). Puedes aprender más acerca de las cGANs a partir de este excelente artículo del blog, y al utilizar este Colab  para probar una cGAN simple.
 
﻿
Run set1
﻿
﻿
﻿
Descripción del Método PropuestoLos trabajos previos que aplicaron transformadores a la generación de imágenes demostraron resultados prometedores para las imágenes con un tamaño de hasta 64x64 píxeles, pero éstas no han podido ser escaladas a una resolución mayor debido al incremento cuadrático de los costos con la longitud de la secuencia. De esta forma, para usar transformadores para sintetizar imágenes de mayor resolución, vamos a necesitar representar a las semánticas de una imagen de forma ingeniosa. Usar la representación de píxeles no va a funcionar, dado que el número de píxeles se incrementa cuadráticamente con un incremento de 2x en la resolución de la imagen. 
Autocodificador Variacional Cuantizado Vectorial (VQ-VAE)Este trabajo está inspirado por el Autocodificador Variacional Cuantizado Vectorial, que difiere de los VAEs tradicionales de dos formas claves:
La red del codificador produce códigos discretos, en vez de continuos (representación latente de la imagen)
La distribución a priori es aprendida, en lugar de la distribución normal multivariante estática.
Puedes aprender más acerca de las VAEs en mi reporte Hacia el Modelado Generativo Profundo con W&B. Pero este trabajo es una evolución más natural de VQ-VAE2 , que mostró lo poderoso que puede ser el aprendizaje de la representación, en el contexto del modelado generativo autorregresivo.
  Figura 2 : Vista general de la arquitectura VQ-VAE. (Fuente)
VQ-VAE consiste de un codificador(E(.)E(.)E(.)﻿), que mapea las observaciones (imágenes) en una secuencia de variables latentes discretas, y de un decodificador(G(.)G(.)G(.)﻿), que reconstruye las observaciones a partir de estas variables discretas. Ellos utilizan un libro de códigos compartido. El libro de códigos está dado por  e∈RK×De \in R^{K \times D}e∈RK×D﻿ Aquí KKK﻿ es el tamaño de los vectores de códigos discretos en el libro de códigos, y DDD﻿  es la dimensionalidad de cada código  eie_iei​﻿, i∈1,2,..,Ki \in 1, 2,..,Ki∈1,2,..,K﻿.
 The quantization is given by, Como es mostrado en la figura 2, se pasa una imagen x a través de EEE﻿ produciendo  E(x)E(x)E(x)﻿. Entonces, esto es cuantizado en base a su distancia a los vectores de código ei, tal que cada vector  E(x)E(x)E(x)﻿ es reemplazado por el índice del vector de código más cercano en el libro de códigos. Lo mismo utiliza el decodificador para la reconstrucción. La cuantización está dada por
			Quantization(E(x))=ekQuantization(E(x)) = e_kQuantization(E(x))=ek​﻿ where k=argmin⁡j∣∣E(x)−ej∣∣k = \underset{j}{\operatorname{argmin}} || E(x) - e_j ||k=jargmin​∣∣E(x)−ej​∣∣﻿﻿
Dos últimas indicaciones rápidas,
La cuantización es un paso no diferenciable, de esta forma, para permitir el entrenamiento constante, el gradiente del error de reconstrucción es propagado hacia atrás a través del decodificador, y hacia el codificador usando el estimador del gradiente directo (copia el gradiente del decodificador al codificador).
Además de la pérdida de la reconstrucción, son utilizadas la pérdida del libro de códigos y la pérdida del compromiso. La pérdida del libro de códigos trae al código seleccionado e cerca de la salida del codificador. La pérdida del compromiso se asegura de que la salida del codificador permanezca cercana al  eee﻿ elegido.
¿Por qué está Funcionando la Representación Latente Discreta? El trabajo Generando Imágenes Diversas de Alta Fidelidad con VQ-VAE-2, que en cierto modo se parece a los Transformadores Adiestrados, utiliza VQ-VAE2 (una modificación a VQ-VAE), ¿pero por qué está funcionando la representación discreta?
Esto está inspirado en la compresión con pérdida JPEG de la imagen. La codificación JPEG elimina más del 80% de los datos, sin cambiar de forma notable la calidad de la imagen percibida. En segundo lugar, entrenar un modelo generativo con menos ruido tiende a funcionar mejor.
Transformador Adiestrado
 Figura 3 :  Diseño de la arquitectura del Adiestramiento de un Transformador 
Para mantener pequeña a la secuencia larga, y para aprovechar la expresividad del transformador, los autores del transformador adiestrado utilizan un libro de códigos discreto de las representaciones aprendidas (inspiradas en VQ-VAE), tal que una imagen  x∈RH×W×3x \in R^{H \times W \times 3}x∈RH×W×3﻿puede ser representada por una colección espacial de las entradas del libro de códigos zq∈Rh×w×nzz_q \in R^{h \times w \times n_z}zq​∈Rh×w×nz​﻿ w, en donde  nzn_znz​﻿  es la dimensionalidad de los códigos.
Como se muestra en la figura 3, los autores han usado VQ-GAN, una variante del VQ-VAE original, que por lo tanto utiliza un discriminador y una pérdida continua para mantener la buena calidad perceptual en tasas de compresión incrementadas.
Este es un diseño de una arquitectura de entrenamiento de dos pasos:
Entrenar a la VQ-GAN y aprender el libro de códigos cuantizados.
Entrenar a un transformador autorregresivo utilizando el libro de códigos cuantizado como una entrada secuencial al transformador.
Entrenando a la VQ-GAN La VQ-GAN es entrenada utilizando un procedimiento de entrenamiento antagónico conocido, con un discriminador D basado en parches. El objetivo completo para encontrar al modelo de compresión óptimo Q∗={E∗,G∗,Z∗}Q^* = \{E^*, G^*, Z^*\}Q∗={E∗,G∗,Z∗}﻿ está dado por
﻿Q∗=argminmax⁡E,G,Z Ex≈p(x)[LVQ(E,G,Z)+λLGAN({E,G,Z},D)]Q^* = \underset{E, G, Z}{\operatorname{argminmax}} \space E_{x \approx p(x)} [L_{VQ}(E, G, Z) + \lambda L_{GAN}(\{E, G, Z\}, D)] Q∗=E,G,Zargminmax​ Ex≈p(x)​[LVQ​(E,G,Z)+λLGAN​({E,G,Z},D)]﻿ 
donde  λ\lambdaλ﻿  es un peso adaptable.
Este procedimiento del entrenamiento reduce significativamente la longitud de la secuencia cuando se desenvuelve el código latente y, de este modo, permite la aplicación de modelos de transformadores poderosos.
Para entrenar tu propia VQ-GAN, puedes clonar el repositorio oficial de GitHub para este paper e instalar las dependencias. Este repositorio ya está instrumentado con Weights and Biases, por lo que puedes obtener automáticamente todas las métricas necesarias en tu tablero de control de W&B. El paso de la preparación de los datos es mencionado con claridad por los autores. Para entrenar la VQ-GAN:
python main.py --base configs/faceshq_vqgan.yaml -t True --gpus 0,
Puedes seleccionar un archivo .yaml diferente, dependiendo del conjunto de datos de tu elección.
Entrenando al transformadorCon el codificador y el decodificador entrenados, ahora podemos representar una imagen en términos de los índices del libro de códigos de sus codificaciones. Esta codificación cuantizada sobre la imagen x está dada por y zq=q(E(x))∈Rh×w×nzz_q = q(E(x)) \in R^{h \times w \times n_z}zq​=q(E(x))∈Rh×w×nz​﻿. Al expresar esta codificación como una secuencia s∈{0,....,∣Z∣−1}h×ws \in \{0,...., |Z|-1\}^{h \times w}s∈{0,....,∣Z∣−1}h×w﻿ de los índices del libro de códigos, que es obtenida al reemplazar cada código por su índice en el libro de códigos  ZZZ﻿, sij=ks_{ij} = ksij​=k﻿  tal que  (zq)ij=zk(z_q)_{ij} = z_k(zq​)ij​=zk​﻿, ahora uno puede mapear la secuencia s de vuelta a la entrada del libro de códigos correspondiente, y puede obtener la imagen decodificada x~=G(zq)\tilde{x} = G(z_q)x~=G(zq​)﻿.
Con esta secuencia, la generación de imágenes puede ser formulada como una predicción autorregresiva del próximo índice. Dados los índices s<is_{<i}s<i​﻿, el transformador aprende a predecir la distribución de los posibles próximos índices, es decir  p(si∣s<i)p(s_i | s_{<i})p(si​∣s<i​)﻿ , para computar la probabilidad de la representación completa como  p(s)=∏ip(si∣s<i)p(s) = \prod_i p(s_i | s_{<i})p(s)=∏i​p(si​∣s<i​)﻿ .De esta forma, la función objetivo es para maximizar la probabilidad logarítmica de las representaciones de los datos. Aquí hay un artículorealmente genial sobre probabilidades.
 Cualquier sistema de generación de imágenes es útil si el usuario puede controlar el proceso de generación. La generación de imágenes puede estar condicionada sobre otra información adicional, como las etiquetas de las clases (como se mostró anteriormente) o las imágenes parciales. Entonces, la tarea es aprender la probabilidad de la secuencia dada esta información ccc﻿:
﻿p(s∣c)=∏ip(si∣s<i,c)p(s|c) = \prod_i p(s_i | s_{<i}, c)p(s∣c)=∏i​p(si​∣s<i​,c)﻿﻿
Con la VQ-GAN entrenada, hay que ajustar la ruta del punto de control de la clave de la configuración model.params.first_stage_config.params.ckpt_path in configs/faceshq_transformer.yaml,  y entonces correr:
python main.py --base configs/faceshq_transformer.yaml -t True --gpus 0,
﻿
﻿
Generando Imágenes de Alta ResoluciónEl costo computacional de entrenar un transformador, que se incrementa cuadráticamente con el incremento de la longitud de la secuencia, pone límites sobre la longitud de la secuencia. De esta manera, para generar imágenes en el marco de los megapíxeles, los autores han tenido que hacer un trabajo de parches y han tenido que recortar imágenes para restringir la longitud de s al máximo tamaño posible durante el entrenamiento.
Para muestrear imágenes, el transformador es utilizado como una ventana deslizante.
Figura 4:  La ventana de atención deslizante
La VQ-GAN se asegura de que el contexto disponible aún sea suficiente como para modelar imágenes con fidelidad, siempre y cuando las estadísticas del conjunto de datos sean, de forma aproximada, espacialmente invariantes o espaciales, condicionando la información que esté disponible.
En el panel mostrado abajo, haz clik en el engranaje y utiliza el control deslizante para visualizar cómo fue sintetizada una imagen de alta resolución utilizando la ventana deslizante.
﻿
﻿
Run set1
﻿
Resultados﻿
Run set1
﻿
﻿
﻿
Run set1
﻿
Lectura Adicional y ConclusiónEl objetivo de este reporte es el de resumir el paper, haciéndolo más accesible a los lectores. He utilizado líneas de este paper en los lugares en donde de esa forma se expresaba mejor la información.
Aquí hay algunas de las lecturas adicionales que podrías encontrar interesantes.
﻿﻿ Hacia el Modelado Generativo Profundo con W&B﻿
﻿﻿ Reflexiones sobre la tipicidad 
﻿Modelos Generativos Autorregresivos (PixelRNN, PixelCNN++)﻿
﻿﻿﻿PixelVAE: Un Modelo Variable Latente para Imágenes Naturales﻿
﻿﻿﻿﻿ Autocodificando un Bit Simple﻿
﻿ ﻿Aprendizaje de la Representación Discreta Neuronal﻿
﻿﻿﻿﻿ Generando Imágenes Diversas de Alta Fidelidad con VQ-VAE-2﻿
﻿