Cómo funciona Stable Diffusion – y Dalle-2 o Midjourney -: Open Source Generative-IA que crea imágenes artísticas (2 de 2)

Ya hemos visto en la primera parte de este artículo cómo funcionan de manera general los modelos de difusión, pero nos queda entender cómo meter en la ecuación el texto que insertamos como entrada. Además, es importante comentar por qué podemos utilizar esta enorme IA como es Stable Diffusion en nuestros ordenadores (aparte de que tenemos disponibilidad de ella al ser Open Source). 

Hacer esto con por ejemplo Dalle-2, si el modelo estuviese disponible para cargar en local, sería imposible ya que Dalle-2 tiene 3.500 millones de parámetros, IMAGEN, 4.600 millones. Stable Diffusion únicamente cuenta con 890 millones, pero su calidad no es menor que el resto de los modelos de IA. Vamos a ver cómo lo consiguieron.

La clave, los espacios latentes

La respuesta es que Stable Diffusion trabaja sobre espacios latentes, sobre codificaciones de muchas menos dimensiones que guardan las características más críticas de las imágenes que recibe (por eso en realidad se lo conoce como un modelo de difusión latente), mientras que el resto lo hace sobre las imágenes en su espacio RGB, con muchas más dimensiones. Estas codificaciones se realizan con el encoder de un Autoencoder Variacional. Al final de todo el proceso, el decoder de este autoencoder se utiliza para transformar la información latente en la imagen RGB de salida.

Esto quiere decir que la U-Net no trabaja sobre las imágenes en su espacio RGB, sino sobre sus versiones comprimidas (latentes), reduciendo en gran medida la carga computacional. Esta era una de las cosas que me quedaba comentar sobre el proceso de denoising (eliminaciones de ruido) con la U-Net, pero falta explicar de qué manera se introduce el texto en este proceso para guiar las generaciones de contenido. Antes, veamos en la imagen siguiente la arquitectura de esta IA.

De izquierda a derecha, recibimos una imagen de entrada RGB y la codificamos en su versión latente (Z) utilizando el encoder (E) del autoencoder, y entonces se empieza el proceso de difusión, es decir, la Cadena de Markov donde se va añadiendo ruido a la representación latente hasta tener la versión final llena de ruido (Zt).

En este momento, se recibe el texto de entrada (conditioning) y comienza el aprendizaje de la U-Net, que recibe la versión latente, el paso de la cadena y el texto para predecir los ruidos e ir eliminándolos, guiando la generación hacia algo que tenga “sentido” con el texto introducido y semejante con lo que ha sido entrenado (es decir, por ejemplo, si usas un dataset de imágenes de Picasso o de Okuda, como en este ejemplo, intentará recrear su estilo), generando con esta información en cada paso versiones menos ruidosas, tantas veces como la longitud de la cadena (T), hasta tener la versión latente final que metemos en el decoder (D) del autoencoder para generar la imagen de salida RGB.

Podemos pensar que, durante el entrenamiento, en el proceso de la Cadena de Markov en el que vamos generando las diferentes versiones latentes ruidosas (Diffusion Process), lo que estamos haciendo es generar una especie de dataset para un aprendizaje supervisado donde cada entrada consiste en tres valores (latente zt, texto, paso de cadena t) y la salida es la cantidad de ruido que se ha generado a partir del latente zt-1 entre el paso de la cadena t-1 y t. 

Como veis, el mismo texto para guiar la salida está presente en cada input del dataset, mientras que en cada uno de este input el latente y el paso de cadena cambia (valores zt y t). De esta manera, vemos como lo que se está haciendo en el denoising es entrenar al modelo a predecir ruidos, en donde al eliminar cada uno de ellos del correspondiente latente se obtendrá una versión menos ruidosa.

La importancia de los Transformers

Habiendo comentado la arquitectura de este modelo de difusión latente, queda responder a la siguiente pregunta: ¿Cómo se integra el texto de entrada en el proceso de denoising? Ha llegado el momento de hablar un poco de Transformers, que son otra de las grandes revoluciones que ocurrieron en el ámbito de la IA, sobre todo en el dominio del procesamiento del lenguaje (¿os suena GPT-3?). 

Como no es el tema central de este artículo, no entraré en detalle en su arquitectura, lo que nos interesa de ella es su componente conocido como Encoder. Por otra parte, tenemos que saber que los modelos de IA son modelos matemáticos, por lo que solo trabajan con números. De esta manera, cada palabra de un texto debe ser codificada antes como un vector de números (Word Embedding) de numerosas dimensiones.

Podemos entender como que cada dimensión del vector significa algo, como por ejemplo “familia”, de esta manera en esta dimensión la palabra “bebé” y “padre” tendrán valores similares. Por lo que también podemos ver que las palabras “cuchara” y “cohete” tendrán vectores de valores muy distintos, ya que no tienen mucho que ver. En la imagen siguiente se muestra un ejemplo, donde por ejemplo la palabra “casas” (houses) está más alejada (tiene menos relación) que el resto de las palabras, que son animales.

Figura 11: Ejemplo de word embedding en 2D

En un espacio de vectores de palabras de muchas dimensiones, existen infinidad de relaciones de este tipo que acaban generando una especie de “diccionario” donde cada palabra tiene su significado y sus relaciones con otras palabras, en forma de vectores numéricos. La clave de las capas de atención que se utilizan en el Encoder de los Transformers es que entrenando este tipo de modelos con enormes conjuntos de datos (textos) se pueden crear unos vectores de palabras muy potentes, ya que almacenan contexto sobre las otras palabras con las que se encuentra en el texto, cosa que ninguna otra arquitectura de IA logra hacer.

Esto que os acabo de explicar es necesario para entender cómo se integran los textos de entrada (vectores numéricos) en el proceso de denoising con la U-Net. Stable Diffusion utiliza los embeddings generados por el encoder de un Transformer perteneciente a una red neuronal (ClipText) pre-entrenada con enormes conjuntos de datos, estos datos con los que se entrena consisten en pares imagen-caption (pie de foto), en el que gracias al entrenamiento aprende a producir embeddings donde por ejemplo una imagen de un gato y la frase “una foto de un gato” están muy relacionados. De esta manera se crea un “diccionario” enorme con vectores de palabras que guardan gran relación con un montón de imágenes de Internet (que es al fin y al cabo de donde se obtienen las imágenes para entrenar el Encoder del Transformer).

Figura 12: Ejemplo pares imagen-caption

De esta manera, en la arquitectura del modelo que vimos en la Figura 9, realmente el texto que se recibe en el proceso de denoising no son las palabras como tal, sino sus codificaciones en vectores numéricos, gracias al uso del Encoder pre-entrenado que acabamos de ver. Una vez hemos visto como llegamos a tener vectores numéricos en vez de las propias palabras que insertamos en el texto, podemos ver ya cómo se integran estos vectores en el proceso de eliminación de ruido. Para que la U-Net soporte el guiado de la generación mediante texto (Text Conditioning) se añaden una serie de capas en su arquitectura denominadas capas de atención.

Como vimos en la Figura 7, la U-Net está formada por una serie de convoluciones que son las que reciben el latente y el paso de la cadena (que, por cierto, este paso no se trata de un número escalar, sino que se codifica también como un vector numérico) y van procesando las características más relevantes. Ahora, después de que la entrada pase por las convoluciones y antes de encogerse en una menor dimensión, se introduce la capa de atención que recibe el vector numérico que representa el texto y lo mezcla con el espacio de información latente generado antes en las convoluciones. Este proceso se repite en las sucesivas convoluciones, de manera que ahora las convoluciones realizadas incorporan información sobre el texto.

¿Cómo se entrenó Stable Diffusion?

Para el entrenamiento, Stable Diffusion se entrenó con el dataset LAION-Aesthetics, consistente de pares imagen-caption, cuya estética de las imágenes se muestran en la imagen siguiente (ésta será la estética que se recreará en las generaciones de imágenes, puesto que el modelo ha sido entrenado con ellas). 

Figura 13: Dataset LAION-Aesthetics

De esta manera, el modelo finalmente va aprendiendo un espacio latente del mundo (de los datos) que ha “visto” y, una vez se tiene el modelo entrenado, introduciendo un texto (que se transformará a vector numérico usando el Encoder del Transformer pre-entrenado) y empezando como semilla con un array latente aleatorio se genera la imagen de salida condicionada a este texto. El modelo analiza todo el “conocimiento” que ha adquirido gracias al entrenamiento para generar la imagen que entiende como “más probable”, “viajando” sobre el espacio latente aprendido. 

Conclusiones

Las nuevas Generative-AI basadas en modelos de difusión han supuesto que mucha más gente haya empezado a hablar sobre la Inteligencia Artificial, pero no solo por gente que se dedica al sector tecnológico, – como hemos aplicado ya en el mundo de la Ciberseguridad – muchas personas ajenas a este sector también han empezado a conocer este campo tan emocionante, como es el de la IA.  

Figura 15: Libro de Machine Learning aplicado a Ciberseguridad de Carmen Torrano, Fran Ramírez, Paloma Recuero, José Torres y Santiago Hernández

Además, un gran sector dedicado al diseño gráfico va a poder aprovecharse de estos modelos para poder integrarlos en sus procesos creativos ya que, por ejemplo, ya existen plugins donde poder usar Stable Diffusion en Blender, herramientas donde con el texto que introduces cambias el material de un objeto, o incluso la gigante Adobe, ya está experimentando con todo esto también. 

¡El futuro es prometedor!