Variational Autoencoders

Antes de empezar con los autoencoders variacionales vamos a hacer un pequeño repaso sobre autoencoders.

Autoencoders

Los autoencoders son una familia de redes neuronales que tienen como objetivo aprender una representación útil 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 representa los datos de entrada $x$ en un espacio latente, $z=g_\phi(x)$. Este espacio suele tener menor dimensión, aunque también existen autoencoders sobrecompletos regularizados.
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 1980 y fue retomada, entre otros trabajos, por Hinton y Salakhutdinov^{[hinton2006reducing]} en 2006. Hoy tiene aplicaciones directas en tareas como 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.

VAE

Los autoencoders deterministas no imponen por sí solos una distribución conocida sobre sus códigos. Pueden dejar regiones del espacio latente sin datos y, si muestreamos arbitrariamente una de ellas, el decoder puede producir una salida poco realista.

Un VAE representa cada entrada mediante una distribución aproximada $q_\phi(z\mid x)$ y regulariza esas distribuciones con respecto a un prior conocido $p(z)$.

La combinación de una posterior continua y la regularización hacia el prior favorece un espacio latente que podemos muestrear: primero tomamos $z\sim p(z)$ y después generamos $x\sim p_\theta(x\mid z)$. Mapear a una distribución, por sí solo, no garantizaría esta propiedad.

Esto es precisamente 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(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(z)\, dz$

En general no es fácil calcular $p_\theta(x)$ porque la integral sobre un espacio latente continuo y de alta dimensión es intratable. Los VAE recurren a la inferencia variacional para aproximar la posterior $p_\theta(z\mid x)$ y construir un objetivo optimizable.

La distribución compleja que queremos aproximar es $p_\theta(z\mid x)$. Elegimos una familia tractable $q_\phi(z\mid x)$ y aprendemos sus parámetros globales $\phi$, que son los pesos del encoder. En el caso habitual de una gaussiana diagonal, el encoder produce para cada entrada una media $\mu_\phi(x)$ y una desviación $\sigma_\phi(x)$.

De esta manera, se introduce:

• Un encoder $q_\phi(z\mid x)$ que aproxima la posterior $p_\theta(z\mid x)$.
• Un decoder $p_\theta(x\mid z)$ que modela la probabilidad de los datos condicionada al código latente.

Para poder retropropagar a través del muestreo se utiliza el truco de reparametrización:

$$\epsilon\sim\mathcal{N}(0,I),\qquad z=\mu_\phi(x)+\sigma_\phi(x)\odot\epsilon.$$

La aleatoriedad queda en $\epsilon$, mientras que $\mu_\phi(x)$ y $\sigma_\phi(x)$ siguen siendo funciones diferenciables de los pesos del encoder.

Función de error: ELBO

La distribución estimada por el encoder, $q_\phi(z\mid x)$, debe aproximar la posterior real $p_\theta(z\mid x)$. Para comparar distribuciones usamos la divergencia de Kullback–Leibler, $D_{\mathrm{KL}}(q\|p)$, que mide el coste esperado adicional de representar muestras de $q$ mediante $p$.

En este caso queremos minimizar $D_{\mathrm{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 el prior.

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

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

conocida como Evidence Lower Bound (ELBO). La ELBO se maximiza; cuando hablamos de función de pérdida minimizamos su negativo, $-\mathcal{L}_{\mathrm{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 minimizando la divergencia KL entre la posterior aproximada y el prior. Normalmente elegimos $p(z)=\mathcal N(0,I)$ y una $q_\phi(z\mid x)$ gaussiana diagonal cuyos parámetros dependen de cada entrada.

Aplicaciones

Las aplicaciones de los VAE incluyen el modelado de texto^{[yang2017improved]}, la detección de anomalías, la estimación de la vida útil de sistemas industriales^{[costa2022variational]}, la clasificación de arritmias^{[costa2021semi]} o el diseño molecular^{[de2023population]}.

Posterior collapse

Los VAE son susceptibles al posterior collapse: el encoder aprende $q_\phi(z\mid x)\approx p(z)$ para todas las entradas, el término KL se aproxima a cero y $z$ deja de contener información útil sobre $x$. Esto ocurre especialmente cuando el decoder es suficientemente expresivo para modelar los datos sin utilizar la variable latente.

No significa que todas las entradas se conviertan necesariamente en un único punto. Significa que sus distribuciones posteriores se vuelven indistinguibles del prior y, por tanto, el decoder puede ignorar $z$. La calidad visible de las muestras dependerá de la capacidad del decoder: pueden seguir siendo plausibles aunque el espacio latente no se utilice.

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.
• Limitar la capacidad del decoder o añadir conexiones desde $z$: reduce la posibilidad de que el decoder explique los datos sin recurrir a la variable latente.

Variantes del VAE

Los VAE entrenados con una pérdida de reconstrucción píxel a píxel tienden a producir imágenes borrosas. ¿De dónde procede este efecto?

Cuando el likelihood de píxeles se modela como una gaussiana con varianza fija, maximizarlo equivale aproximadamente a minimizar un error cuadrático. Si una entrada admite varias reconstrucciones plausibles, ese objetivo favorece la media condicional de esas alternativas, que puede verse borrosa. No se trata de promediar las reconstrucciones de todo el dataset, sino de una consecuencia del modelo de observación y de la incertidumbre multimodal.

β-VAE

Una familia de objetivos introduce un peso $\beta$ sobre la regularización:

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

El β-VAE original^{[higgins2017beta]} utiliza habitualmente $\beta>1$ para favorecer representaciones más disentangled a costa de fidelidad de reconstrucción. El objetivo general también permite explorar otros valores.

Si $\beta=1$ recuperamos la ELBO estándar. Con $0<\beta<1$ se penaliza menos la información latente y puede mejorar la reconstrucción, aunque la posterior agregada puede alejarse del prior. Con $\beta>1$ aumenta la presión de regularización y compresión.

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 la reconstrucción se vea bien" y "quiero que el espacio latente tenga sentido".

InfoVAE

El InfoVAE^{[zhao2017infovae]} separa el control de la información mutua de la coincidencia entre la posterior agregada $q_\phi(z)$ y el prior. Una forma de su objetivo es:

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

Aquí $D$ puede ser KL, MMD u otra divergencia adecuada. Los hiperparámetros $\alpha$ y $\lambda$ controlan la información conservada sobre $x$ y cuánto se aproxima la posterior agregada al prior; determinados valores recuperan el VAE estándar.

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 se recopilan algunos de los trabajos más influyentes.

Hierarchical VAEs

En el VAE básico se utiliza una sola variable latente con un prior sencillo, normalmente una gaussiana estándar. Una jerarquía introduce variables a distintas escalas y priors condicionales, lo que aumenta la expresividad tanto del modelo de inferencia (encoder) como del modelo generativo (decoder) y puede mejorar la log-verosimilitud:

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:

$\mathcal{L}_{\mathrm{ELBO}} = \mathbb{E}_{q(z_{1:k}\mid x)}[\log p(x \mid z_1)] - D_{\mathrm{KL}}\!\left(q(z_{1:k}\mid x)\,\|\,p(z_{1:k})\right)$

Se han publicado muchos modelos jerárquicos, 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.

VD-VAE muestra hasta dónde puede llegar una jerarquía muy profunda construida con bloques convolucionales y distribuciones gaussianas diagonales. Utiliza nearest-neighbor upsampling seguido de convoluciones, parametrizaciones residuales y otras decisiones de estabilidad para entrenar directamente con la ELBO. El remuestreo evita artefactos asociados a ciertas convoluciones transpuestas, pero no debe interpretarse por sí solo como una solución general al posterior collapse.

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. Vamos a verlo, 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\mid y),K(y\mid x)\}=K(x,y)-\min\{K(x),K(y)\}+O(\log K)$$

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 descompresió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 de codificación entrópica, como Huffman, asignan códigos prefijo más cortos a los símbolos más probables ^{[moffat2019huffman]}.

Los avances recientes han incorporado técnicas de Machine Learning a los algoritmos de compresión. ¿Qué modelos podemos 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$

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

Los métodos de compresión sin pérdida buscan minimizar el número de bits sin sacrificar información. Para un símbolo de probabilidad $p(x)$, la longitud ideal es $-\log_2 p(x)$ bits; la longitud media óptima está acotada por la entropía de Shannon.

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 corresponde inicialmente a $-\log_2 p(z)-\log_2 p(x\mid z)$. El método bits-back permite recuperar aproximadamente $-\log_2 q(z\mid x)$ bits usados para seleccionar $z$^{[frey1997efficient]}. Esta idea conecta la inferencia variacional con la compresión probabilística.

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_{\mathrm{extra}}+\log_2 q(z\mid x)-\log_2 p(x\mid z)-\log_2 p(z)$$

Los términos $-\log_2 p(z)$ y $-\log_2 p(x\mid z)$ se transmiten, mientras que $-\log_2 q(z\mid x)$ bits se recuperan mediante bits-back. En esperanza, el coste neto coincide con la negative Evidence Lower Bound (negative ELBO):

$$\mathbb E_{q(z\mid x)}[N-n_{\mathrm{extra}}] = \mathbb E_{q(z\mid x)}\left[\log_2\frac{q(z\mid x)}{p(x,z)}\right] =-\mathcal L_{\mathrm{ELBO}}/\log 2$$

La división por $\log 2$ convierte nats en bits. Este coste corresponde a la función de pérdida que minimizamos en un VAE.

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

• El primer término corresponde a la reconstrucción de $x$ y favorece un esquema de codificación y decodificación eficiente maximizando $\log p(x\mid z)$ con muestras de $q(z\mid 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 la negative ELBO maximiza un límite inferior de $\log p(x)$ y aproxima $q(z\mid x)$ a la posterior del modelo $p(z\mid 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.