VAE

Antes de meternos de lleno con los VAEs vamos a hacer un pequeño repaso sobre autoencoders.

Autoencoders

Los autoencoders son una familia de redes neuronales que tiene como objetivo aprender una representación simplificada de los datos de entrada.

Su funcionamiento se basa en comprimir y reconstruir datos. Generalmente, constan de dos redes:

1. Un Encoder $g(\cdot)$ que comprime los datos de entrada $x$ a un espacio de menores dimensiones, conocido como espacio latente, $z = g_\phi(x)$.
2. Un Decoder $f(\cdot)$ que a partir del espacio latente reconstruye los datos, $x' = f_\theta(g_\phi(x))$.

Los parámetros de cada red $(\phi, \theta)$ se optimizan a la vez para que las reconstrucciones sean lo más parecidas a los datos originales, $x \approx f_\theta(g_\phi(x))$.

La idea se originó en la década de los 80 y posteriormente fue promovida por Geoffrey Hinton y compañía^{[hinton2006reducing]} en 2006. Es una idea muy intrínseca a los modelos de IA actuales y tiene una aplicación directa a diversas tareas como la detección de anomalías, manifold learning, compresión o generación de datos.

La función de pérdida de un autoencoder convencional mide qué tan bien el modelo puede reconstruir la entrada original una vez comprimidos los datos al espacio latente, por tanto, el objetivo es que la salida $\hat{x}$ sea lo más parecida posible a la entrada $x$. La elección de la función de error depende principalmente de la naturaleza de los datos de entrada.

• Mean Squared Error (MSE)

  Es la función de pérdida más utilizada para datos continuos (por ejemplo, imágenes o series temporales).

  El objetivo es minimizar la diferencia al cuadrado entre la entrada original $x$ y la salida reconstruida $\hat{x}$:

  $MSE = \frac{1}{n} \sum_{i=1}^{n} (x_i - \hat{x}_i)^2$

• Binary Cross Entropy

  Se utiliza cuando los datos de entrada están normalizados en el rango [0, 1] y se interpretan como probabilidades o cuando los datos son binarios. Es muy común en autoencoders donde la capa de salida tiene una función de activación sigmoide.

  $BCE = -\frac{1}{n} \sum_{i=1}^{n} [x_i \log(\hat{x}_i) + (1 - x_i) \log(1 - \hat{x}_i)]$

Existen muchas variantes de autoencoders como Denoising Autoencoders^{[vincent2008extracting]}, Sparse Autoencoders^{[makhzani2013k]} o Contractive Autoencoders^{[rifai2011contractive]}. De aquí en adelante nos centraremos en los Variational AutoEncoders (VAE)^{[kingma2013auto]} por sus propiedades generativas.

En esta librería tenéis varios modelos de autoencoders programados en Keras 3. Se agradece una estrella 😊.

VAE

Los autoencoders estándar traducen las entradas a un conjunto de vectores fijo y esto ocasiona que el espacio latente sea discontinuo. Esto se traduce a que si se quisiera elegir un punto del espacio latente donde no existe representación alguna y samplear con el decoder para generar una nueva muestra, la salida será poco realista, porque el decoder no sabe cómo reconstruir una entrada a partir de una región del espacio latente vacía. Con 'vacía' me refiero a que durante el entrenamiento el modelo nunca procesó datos que se codificaran en ese punto del espacio latente.

La novedad que introducen los VAE es que en lugar de mapear los datos a vectores fijos se mapean a una distribución.

Esto tiene la ventaja de que el espacio latente es continuo y, por tanto, se puede elegir un punto $z$ que no pertenezca a la proyección de ningún dato de entrenamiento, pasarlo por el decoder y obtener una nueva muestra $x$ coherente. Es decir, generar datos en la forma $z \rightarrow x$.

Ya tenemos un modelo generativo de variables latentes!

La relación entre los datos de entrada $x$ y el vector latente $z$ viene dada por:

• Prior: $p_theta(z)$
• Likelihood: $p_\theta(x \mid z)$
• Posterior: $p_\theta(z \mid x)$

La distribución de los datos entonces se puede modelar como:

$p_\theta(x) = \int p_\theta(x \mid z), p_\theta(z), dz.$

Desafortunadamente, no es fácil calcular $p_\theta(x)$, ya que es muy costoso comprobar y sumar todos los valores posibles de $z$ (por lo que se suele denominar 'intratable'). Para solucionar este problema, los VAE recurren a la inferencia variacional (de ahí su nombre). Precisamente, la inferencia variacional es uno de los métodos más recurridos de la inferencia bayesiana y se utiliza para aproximar integrales intratables. En otras palabras, es una técnica utilizada para aproximar distribuciones complejas.

La distribución compleja que se quiere aproximar es $p_\theta(z \mid x)$ (posterior). Para aproximarla mediante inferencia variacional, la idea es presuponer una familia de distribuciones simples $q_\phi$ y plantear la aproximación como un problema de optimización: buscar dentro de la familia de distribuciones $q_\phi$ la distribución que más se aproxime a la distribución objetivo. $\phi$ denota los parámetros de la familia de distribuciones elegida. Por ejemplo, si asumimos que va a ser una gaussiana, $\phi$ serían la media y la varianza.

De esta manera, se introduce:

• Una función de aproximación $q(z \mid x)$, modelada mediante un encoder para aprender esos parámetros $\phi$. Es decir, un mecanismo para aproximar la posterior $p(z \mid x)$.
• Un decoder que se encarga de modelar la probabilidad condicional $p(x \mid z)$. Al aprender $\theta$, está aprendiendo cómo mapear cualquier punto $z$ del espacio latente de vuelta al espacio original de los datos.

Función de error: ELBO

La distribución estimada por el encoder, $q_\phi(z \mid x)$, debe ser próxima a la real $p_\theta(z \mid x)$. Para cuantificar la distancia entre estas dos distribuciones se usa la divergencia de Kullback–Leibler, $D_{\text{KL}}(Y \| X)$, que se encarga de medir cuánta información se pierde si se usa la distribución $Y$ para representar $X$.

En este caso, se quiere minimizar $D_{\text{KL}}(q_\phi(z \mid x) ,\|\, p_\theta(z \mid x))$ con respecto a $\phi$.

Usando el teorema de Bayes, se puede reescribir como:

$$D_{\text{KL}}(q_\phi(z \mid x) \,\|\, p(z \mid x)) = \log p(x) - \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)] + D_{\text{KL}}(q_\phi(z \mid x) \,\|\, p(z))$$

• El primer término es el log-likelihood de los datos, que queremos maximizar.
• El segundo término es el log-likelihood esperado de los datos bajo la posterior aproximada (encoder).
• El tercer término es la divergencia KL entre la posterior aproximada y la prior.

Combinando estos términos, se puede definir la función de pérdida de un VAE como:

$\mathcal{L}(\theta, \phi; x) = \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)] - D_{\text{KL}}(q_\phi(z \mid x),\|\, p_\theta(z)),$

también conocida como Evidence Lower Bound (ELBO).

• El primer término es la reconstrucción de $x$ que tiende a hacer el esquema de codificación-decodificación lo más eficiente posible maximizando el log-likelihood $\log p_\theta(x \mid z)$ con muestreo de $q_\phi(z \mid x)$ (encoder).
• El segundo término regulariza las variables latentes (representadas por $z$) minimizando la divergencia KL entre la aproximación variacional (encoder) y la distribución a priori de $z$. Normalmente, la familia de distribuciones simples $q_\phi$ que elegimos es una normal estándar, por tanto, cuando pasas una entrada $x$ por el encoder este te devuelve los valores de $\mu$ y $\sigma$ para esa entrada específica.

Aplicaciones

Las aplicaciones de los VAE son innumerables. Algunos ejemplos más allá de la generación de imágenes son su aplicación a modelado de texto y su semántica^{[yang2017improved]}, detección de anomalías^{[lin2020anomaly]}, estimación de la vida útil de sistemas industriales^{[costa2022variational]}, detección de arritmias^{[costa2021semi]} o diseño molecular^{[de2023population]}.

Posterior collapse

Los VAE son conocidos por ser susceptibles al posterior collapse. Esto es un problema que se refiere a la situación en la que la posterior aproximada por el encoder colapsa con la prior. Esto signfica que la KL $\rightarrow$ el segundo término de la ELBO se vuelve dominante. Esto causa que el término de reconstrucción tenga muy poco peso y que, por tanto, el decoder tienda a reconstruir los datos no a partir de $q(z \mid x)$ si no a partir de $p(z)$. Es decir, no tiene en cuenta los datos!!

En consecuencia, el modelo decide que la mejor forma de minimizar el error es haciendo que todas las entradas colapsen en un solo punto en el espacio latente $\rightarrow$ es la representación más fiel a una normal estándar $p(z)$, pero nosotros queremos $q(z \mid x)$.

En este caso, las reconstrucciones/generaciones del decoder son incoherentes.

Trucos de entrenamiento

Hay varias aproximaciones para solventar este problema, entre ellas:

• KL Divergence Annealing: incrementa gradualmente el peso del término de divergencia KL en la función de pérdida durante el entrenamiento para permitir que el modelo aprenda primero a reconstruir los datos antes de imponer una regularización fuerte en el espacio latente.
• Free Bits: establece un umbral mínimo para la divergencia KL en cada dimensión de la variable latente, asegurando que cada dimensión contribuya con al menos un "mínimo de información" y evitando que colapsen a valores triviales.
• Minimum Desired Rate: similar a Free Bits, pero aplica una penalización para mantener un mínimo de información en el espacio latente, forzando que la compresión no pierda relevancia.
• Dropout: desconecta aleatoriamente conexiones de la red durante el entrenamiento para evitar que el modelo dependa demasiado de ciertas neuronas, manteniendo así la relevancia de las variables latentes.
• Independent Hidden States: obliga a que las representaciones ocultas del encoder sean independientes entre sí, asegurando que las variables latentes no se vuelvan redundantes o que algunas dimensiones se ignoren.

Variantes del VAE

La literatura de los VAEs es muy extensa y se han propuesto múltiples variantes, así que comentaré solo algunas de ellas.

Otro problema en el entrenamiento de los VAE es que generan imágenes borrosas. Sin embargo, esto no pasa con modelos que optimizan explícitamente la probabilidad de los datos, como los normalizing flows, entonces... ¿por qué aquí sí que ocurre?

Normalmente para cuantificar la calidad de las reconstrucciones se utilizan métricas como RMSE o BCE, como comentamos antes; por lo tanto, la red tiene como objetivo minimizar estas métricas.

Durante el entrenamiento, el modelo puede acertar una reconstrucción en concreto pero fallar en el resto, por lo que el error será alto. Entonces las redes hacen la siguiente trampa: en lugar de acertar reconstrucciones concretas producen para todos los datos generaciones borrosas pero lo suficientemente cercanas a los datos reales, de forma que cuando se hace la media de las reconstrucciones el error obtenido sea bajo.

β-VAE

Se puede resolver este problema modificando el encoder para evitar fusionar entradas distintas en el mismo espacio latente, o el decoder añadiendo información faltante. Sin embargo, una solución aún más sencilla es reducir la penalización sobre el término KL, haciendo que el modelo se acerque más a un autoencoder determinista:

$\mathcal{L}(\theta, \phi; x) = \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)] - \beta\, D_{\text{KL}}(q_\phi(z \mid x) \,\|\, p_\theta(z)).$

En esto consiste el β-VAE^{[higgins2017beta]} (2017).

Si se fija $\beta = 1$, la función objetivo es la misma utilizada en los VAE estándar; si se fija a 0, entonces es la misma utilizada por los autoencoders convencionales. Si se usa $\beta < 1$, se almacenan más bits sobre cada entrada (es decir, se da más importancia a los datos) y, por tanto, se pueden reconstruir las imágenes de forma menos borrosa. Si se utiliza $\beta > 1$, se obtiene entonces una representación más comprimida.

Con $\beta < 1$ se puede reducir en parte las generaciones borrosas. En contrapartida, una ventaja de utilizar $\beta > 1$ es que fomenta el aprendizaje de una representación latente "desenredada". Intuitivamente, esto significa que cada dimensión latente representa un factor/característica diferente de los datos de entrada. Por ejemplo, un modelo entrenado con fotos de rostros humanos podría captar la suavidad, el color de la piel, el color o la longitud del pelo, la emoción, si se llevan gafas y muchos otros factores relativamente independientes en dimensiones separadas.

En resumen: el diseño de un VAE es casi siempre un juego de equilibrio entre "quiero que se vea bien" y "quiero que el espacio latente tenga sentido".

InfoVAE

Para solucionar el problema de poner demasiado énfasis en el segundo término (la regularización), causando representaciones latentes que no capturan bien la estructura de los datos y muestras generadas que no son diversas, el InfoVAE^{[zhao2017infovae]} (2017) modifica la función de pérdida añadiendo un término de maximización de la información mutua entre las variables latentes $z$ y los datos observados $x$. La función de pérdida de InfoVAE se puede escribir como:

$\mathcal{L}_{\text{InfoVAE}} = \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)] - \alpha\, D_{\text{KL}}(q_\phi(z \mid x) \,\|\, p_\theta(z)) + \lambda\, D_{\text{KL}}(q_\phi(z) \,\|\, p_\theta(z)).$

Donde el primer término es el error de reconstrucción, el segundo es la divergencia KL entre la distribución posterior aproximada y la prior, y el tercero es la divergencia KL entre la distribución marginal de las variables latentes y la distribución prior.

$\alpha$ y $\lambda$ son hiperparámetros que controlan la importancia relativa de cada término. De esta manera, al maximizar la información mutua entre $z$ y $x$, se asegura que las representaciones latentes $z$ contengan más información relevante sobre los datos $x$.

Multimodal VAEs

Los VAE multimodales^{[wu2018multimodal, shi2019variational]} son una extensión del modelo VAE estándar para manejar y aprender representaciones conjuntas de datos provenientes de diferentes modalidades. Estas modalidades pueden incluir texto, imágenes, audio, vídeo, etc. La idea principal es capturar las dependencias y correlaciones entre diferentes tipos de datos, aprendiendo una representación latente compartida que pueda ser utilizada para tareas como la generación o imputación de datos.

Hasta ahora se han propuesto varios enfoques para el aprendizaje multimodal de VAEs; en esta review^{[sejnova2023benchmarkingmultimodalvariationalautoencoders]} se recopilan algunos de los trabajos más influyentes.

Hierarchical VAEs

Un problema de los VAE es que sufren de priors simples. Como hemos visto, asumimos que la prior es una Normal estándar. Una forma de suavizar esto es con jerarquías de variables latentes. Jerarquizando tanto el modelo de inferencia (encoder) como el modelo generativo (decoder) se consiguen mejores probabilidades:

La distribución que se quiere modelar entonces es:

$p(x) = \int p(x \mid z_1)\, p(z_1 \mid z_2) \cdots p(z_k)\, dz_{1:k},$

Por tanto, el ELBO se expresaría de la siguiente forma:

$\text{ELBO} = \mathbb{E}_{q(\cdot \mid x)}[\log p(x \mid z_{1:k})] - \text{KL}\!\left( q(z_{1:k} \mid x) \,\|\, p(z) \right).$

Se han publicado muchos trabajos que exploran diferentes tipos de modelos HVAE, entre los que destacan VD-VAE (Very Deep VAE)^{[child2020very]} para generación de imágenes o Bit-swap^{[kingma2019bit]} para compresión de datos.

La arquitectura de VD-VAE es un VAE convolucional simple. Para cada capa, el prior y el posterior son gaussianas diagonales. Descubrieron una técnica en la que el remuestreo del vecino más cercano (en el decoder) funcionaba mucho mejor que la convolución transpuesta y evitaba el posterior collapse. Esto permitió el entrenamiento con el ELBO original, sin necesidad de ninguno de los trucos discutidos anteriormente.

Extra: Compresión

En Bit-swap se presenta un VAE jerárquico con el que se demuestra que el esquema de compresión bits-back (BB-ANS) se puede usar con modelos de variables latentes. Es decir, este tipo de modelos generativos se pueden utilizar como esquemas de compresión eficientes. Veámoslo, pero primero: un poco de teoría de la información.

Complejidad de Kolmogórov

En teoría de la información, la complejidad de Kolmogórov de un objeto $x$, denotada como $K(x)$, se refiere a la longitud de la descripción más corta posible de dicho objeto en un 'ordenador universal'. Representa la complejidad inherente de una pieza de información $x$ en función de la longitud del programa binario más conciso capaz de generarla.

$\rightarrow$ Un programa binario más corto implica un mayor nivel de regularidad o predictibilidad en los datos.

$\rightarrow$ Un programa binario más largo sugiere una estructura más compleja y menos compresible.

De manera similar, la complejidad de Kolmogórov de $x$ dado $y$, $K(x|y)$, corresponde a la longitud del programa binario que, recibiendo $y$ como entrada, produce $x$. A partir de esta idea, Bennett et al. ^{[bennett1998information]} derivaron la métrica de distancia de información, $E(x,y)$:

$$E(x, y) = \max(K(x|y),K(y|x))= K(xy)-min(K(x), K(y))$$

Esta métrica evalúa esencialmente la existencia de un programa sencillo capaz de transformar un objeto en otro y, por tanto, proporciona una medida de la relación o de la cantidad de información compartida entre ellos. Cuanto más simple sea el programa de conversión, más similares se consideran los objetos. Sin embargo, la complejidad de Kolmogórov no es computable teóricamente sino que solo puede aproximarse.

Compresión

Las técnicas de compresión aproximan la complejidad de Kolmogórov, ya que ofrecen un medio práctico para representar los datos de manera más concisa sin perder información esencial. El objetivo principal de un compresor sin pérdida (lossless compressor) es minimizar el número de bits necesarios para representar los datos, permitiendo posteriormente su decompresión exacta.

Los métodos de compresión logran este objetivo identificando y explotando redundancias o patrones presentes en los datos. Para alcanzar la menor longitud de compresión posible, los símbolos con mayor probabilidad reciben códigos más cortos. De acuerdo con el Teorema de Codificación de Fuentes de Shannon ^{[shannon1948mathematical]}, la longitud de los bits codificados está directamente relacionada con la entropía de la fuente de información:

Lo que buscan los algoritmos de compresión es modelar eficazmente la distribución “real” de los datos $p(x)$:

• Una estrategia habitual consiste en los métodos basados en diccionarios, donde se identifican patrones o secuencias de datos y se reemplazan por códigos más cortos. Entre ellos se encuentran los algoritmos de Lempel-Ziv, que construyen dinámicamente un diccionario durante la compresión ^{[kosaraju2000compression]}.

• Los modelos basados en contexto, como la codificación aritmética adaptativa, ajustan las probabilidades de codificación en función del contexto de los símbolos previamente codificados ^{[moffat1998arithmetic]}.

• Los métodos basados en transformaciones, como la ampliamente utilizada codificación de Huffman, reorganizan los datos en una forma más fácilmente compresible transformando la representación de los símbolos ^{[moffat2019huffman]}.

Avances más recientes han incorporado técnicas de Machine Learning en los algoritmos de compresión, y qué tipo de modelos se puede utilizar para aproximar la distribución real de los datos?

$\rightarrow$ Exacto, los VAEs!!

Modelos generativos como compresores

En un pipeline habitual de compresión, los datos se representan mediante una secuencia de símbolos $X = x_1, x_2, ..., x_n$ y se codifican en un conjunto de bits $C$. Posteriormente, se aplica el proceso inverso para recuperar los datos originales.

• Compresión: $E(X) \rightarrow C$

• Decompresión: $D(C) \rightarrow X'$

Los métodos de compresión sin pérdida buscan minimizar el número de bits necesarios para representar un archivo sin sacrificar información. Los símbolos pertenecen a un conjunto finito y aparecen con probabilidades conocidas. La tasa óptima de compresión viene dada por la entropía de Shannon, donde $p(x)$ representa la distribución de los datos y $log(p(x))$ aproxima el número de bits requeridos para codificarlos.

Cuanto mejor sea la compresión, mejor será la aproximación de $p(x)$. En un contexto probabilístico, $p(x)$ puede expresarse como la probabilidad de observar los datos bajo un modelo determinado:

$$p(x) = \int p(x|\theta) p(\theta), d\theta$$

donde $p(x|\theta)$ es la probabilidad de observar los datos dado un conjunto específico de parámetros, $p(\theta)$ es la distribución a priori y $\theta$ representa los parámetros del modelo.

En el contexto de un modelado generativo, esta expresión suele formularse en términos de la variable latente $z$:

$$p(x) = \int p(x|z), p(z), dz$$

Te suena, no?

Utilizando esta notación, el proceso de envío de información en un esquema de compresión ($E(X)\rightarrow C$) se desarrolla de la siguiente manera:

1. La variable latente $z$ se codifica utilizando $p(z)$.
2. La muestra $x$ se codifica utilizando $p(x|z)$.
3. Se transmiten tanto $z$ como $x$.

Por otra parte, la recepción de la información ($D(C)\rightarrow X'$) sigue estos pasos:

1. $z$ se decodifica utilizando $p(z)$.
2. $x$ se decodifica utilizando $p(x|z)$.

El número de bits involucrados en este proceso corresponde a la suma de los bits necesarios para transmitir $z$ y los necesarios para transmitir $x$: $-log(p(z))$ y $-log(p(x|z))$. Sin embargo, es posible transmitir algunos bits adicionales sin coste efectivo, concretamente $log(q(z|x))$, tal como describe el método bits-back ^{[frey1997efficient]}. Este concepto mantiene una estrecha relación con la inferencia variacional y los VAEs, como se ha mostrado en diversos trabajos previos ^{[honkela2004variational]}, ^{[chen2016variational]}.

El proceso de compresión descrito anteriormente puede interpretarse según el esquema BB-ANS ^{[townsend2019practical]}. Algunos bits adicionales se utilizan para generar una muestra $z$ mediante el encoder $q(z|x)$; posteriormente, $z$ se emplea para codificar $x$ a través de $p(x|z)$, mientras que la propia variable $z$ se codifica usando la distribución a priori $p(z)$. El proceso inverso permite recuperar $z$ mediante $p(z)$, reconstruir $x$ mediante $p(x|z)$ y recuperar los bits adicionales iniciales utilizando $q(z|x)$.

La longitud final del flujo de bits para un único dato viene dada por:

$$N = n_{extra} + log(q(z|x)) - log(p(x|z)) - log(p(z))$$

Fíjate que los términos $-log(p(z))$ y $-log(p(x|z))$ siguen transmitiéndose, pero los $log(q(z|x))$ bits pueden recuperarse mediante el mecanismo bits-back. Lo relevante aquí es que $N - n_{extra}$ coincide con el negativo de la (nEvidence Lower Bound, nELBO):

$$E_{q(z|x)} [N-n_{extra}] = E_{q(z|x)} log\frac{p(x,z)}{q(z|x)} = -ELBO$$

que corresponde precisamente con la función de pérdida utilizada en los VAEs!!

$$ELBO = E_{q(z|x)} \left[log\frac{p(x,z)}{q(z|x)}\right] = E_{q(z|x)}\left[log\frac{p(x|z)p(z)}{q(z|x)}\right] = E_{q(z|x)}\left[log p(x|z)\right] + E_{q(z|x)}\left[log\frac{p(z)}{q(z|x)}\right] = E_{q(z|x)}\left[log p(x|z)\right]-D_{KL}\left[q(z|x)||p(z)\right]$$

• El primer término corresponde a la reconstrucción de $x$ y favorece a que el esquema de codificación/decodificación sea lo más eficiente posible maximizando $log p(x|z)$ utilizando muestras provenientes de $q(z|x)$.

• El segundo término es la regularización del espacio latente mediante la minimización de la divergencia KL entre la aproximación variacional y la distribución a priori $z$.

En consecuencia, minimizar el nELBO equivale simultáneamente a maximizar la probabilidad de generar datos reales, $log p(x)$, y minimizar la divergencia KL entre $q(z|x)$ y $p(z|x)$.

Esta equivalencia demuestra que un modelo de variables latentes correctamente optimizado, como un VAE, es directamente aplicable a tareas de compresión, ya que minimiza la longitud de código alcanzable en el contexto de la codificación bits-back.

Jiang et al., en NPC (Non-Parametric learning by Compression)^{[jiang2022few]}, propusieron un método basado en esta premisa para clasificación de imágenes, superando a métodos supervisados y semisupervisados.

De igual forma, en Few-shot generative compression approach for system health monitoring ^{[costa2025few]} propuse un método basado en compresión y VAEs para diagnosticar casos donde los datos de monitorización son abundantes, pero hay muy pocas instancias etiquetadas.