Autocodificadores variacionales (VAE): qué son, cómo funcionan y cómo se construyen

Un autocodificador variacional (VAE) no aprende a comprimir una entrada en un único punto: aprende a comprimirla en una distribución de probabilidad. Esa diferencia, que parece un detalle de implementación, es la que separa a un autocodificador que solo sabe reconstruir de un modelo que además sabe generar datos nuevos y coherentes. Este artículo explica qué es un VAE, cómo funciona su mecánica interna (encoder probabilístico, truco de reparametrización, función de pérdida ELBO) y cómo se construye uno de principio a fin con TensorFlow y Keras, entrenado sobre MNIST y con una visualización del espacio latente que hace visible por qué funciona todo lo anterior.

De autocodificador a autocodificador variacional

Un autocodificador tradicional es una red no supervisada con dos partes: un encoder que mapea la entrada a un espacio de menor dimensión (el espacio latente) y un decoder que reconstruye la entrada original a partir de ese espacio. El entrenamiento minimiza la diferencia entre la entrada y su reconstrucción, y el resultado es excelente para reducir dimensionalidad o detectar anomalías.

Lo que un autocodificador tradicional no garantiza es que ese espacio latente sirva para generar nada nuevo. El encoder aprende a mapear cada entrada a un punto, pero no hay ninguna razón para que los puntos cercanos entre sí en ese espacio correspondan a salidas parecidas: si tomas un punto aleatorio y lo pasas por el decoder, no hay ninguna garantía de que el resultado sea coherente.

Un VAE resuelve exactamente ese problema cambiando qué produce el encoder. En lugar de un vector z único por entrada, el encoder produce los parámetros de una distribución gaussiana: una media mu y un logaritmo de varianza log_var. Para generar la representación latente, se muestrea z de N(mu, sigma^2). Como esa distribución se regulariza durante el entrenamiento para acercarse a una gaussiana estándar conocida, el espacio latente completo queda estructurado y es continuo: cualquier punto que muestrees produce una salida razonable, y puntos cercanos producen salidas parecidas.

Los componentes: encoder, decoder y el truco de reparametrización

El encoder de un VAE (también llamado red de inferencia) no produce z directamente: produce mu y log_var. Se trabaja con el logaritmo de la varianza en lugar de la varianza directamente por dos razones prácticas: garantiza que la varianza real (exp(log_var)) sea siempre positiva sin restringir la salida de la red, y mejora la estabilidad numérica del entrenamiento.

El decoder (la red generativa) hace lo mismo que en un autocodificador normal: toma una muestra z del espacio latente y la transforma de vuelta al espacio de datos original.

El problema es que el muestreo de z es una operación aleatoria, y el descenso de gradiente no puede propagarse a través de una operación no diferenciable. El truco de reparametrización lo resuelve moviendo la aleatoriedad fuera del grafo de cómputo: en lugar de muestrear z directamente de N(mu, sigma^2), se muestrea un valor auxiliar epsilon de una normal estándar N(0, 1) y se calcula:

z = mu + exp(0.5 * log_var) * epsilon

Ahora mu y log_var son funciones diferenciables de los pesos del encoder, epsilon queda fuera de ese camino, y el gradiente fluye con normalidad durante el entrenamiento.

La función de pérdida: ELBO, reconstrucción y divergencia KL

El objetivo de entrenamiento de un VAE es maximizar una cota inferior de la evidencia (ELBO, Evidence Lower Bound), formulada originalmente en Auto-Encoding Variational Bayes de Kingma y Welling. En la práctica, minimizar la pérdida equivale a maximizar esa cota, y la pérdida se descompone en dos términos que se suman:

• Pérdida de reconstrucción: mide qué tan bien el decoder reconstruye la entrada x a partir de z. Se usa entropía cruzada binaria para datos en [0,1] (como dígitos MNIST normalizados) o error cuadrático medio para datos continuos.
• Divergencia KL: distancia entre la distribución que aprende el encoder, q(z|x) = N(mu, sigma^2), y el prior p(z) = N(0, I). Para dos gaussianas tiene forma cerrada: KL = -0.5 * sum(1 + log_var - mu^2 - exp(log_var)). Minimizarla es lo que empuja al espacio latente a ser continuo y estructurado en lugar de un conjunto de puntos aislados.

Sin el término KL, el encoder podría simplemente aprender a mapear cada entrada a un punto fijo con varianza casi nula, comportándose como un autocodificador determinista y perdiendo la capacidad de generar muestras nuevas y variadas. El equilibrio entre ambos términos es, como se ve más abajo, también la fuente del fallo de entrenamiento más común en VAEs.

Cómo se construye un VAE con TensorFlow y Keras

La implementación siguiente sigue el mismo patrón que el ejemplo oficial de VAE en la documentación de Keras (capa de muestreo personalizada, train_step propio, entrenamiento sobre MNIST) y se ejecuta tal cual con pip install tensorflow matplotlib numpy instalado.

1. Datos y parámetros

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
import matplotlib.pyplot as plt

# Concatenamos train y test: para este ejercicio no evaluamos generalizacion,
# solo queremos aprender un espacio latente sobre el conjunto completo de digitos.
(x_train, _), (x_test, _) = keras.datasets.mnist.load_data()
mnist_digits = np.concatenate([x_train, x_test], axis=0)
mnist_digits = np.expand_dims(mnist_digits, -1).astype("float32") / 255

original_dim = 28 * 28
intermediate_dim = 256
latent_dim = 2  # en 2 dimensiones, para poder dibujar el espacio latente completo

2. La capa de muestreo (Sampling)

Esta capa encapsula exactamente el truco de reparametrización descrito antes:

class Sampling(layers.Layer):
    """Usa (z_mean, z_log_var) para muestrear z, el vector latente."""
    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = tf.shape(z_mean)[0]
        dim = tf.shape(z_mean)[1]
        epsilon = tf.random.normal(shape=(batch, dim))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon

3. El encoder

encoder_inputs = keras.Input(shape=(28, 28, 1))
x = layers.Flatten()(encoder_inputs)
x = layers.Dense(intermediate_dim, activation="relu")(x)
z_mean = layers.Dense(latent_dim, name="z_mean")(x)
z_log_var = layers.Dense(latent_dim, name="z_log_var")(x)
z = Sampling()([z_mean, z_log_var])
encoder = keras.Model(encoder_inputs, [z_mean, z_log_var, z], name="encoder")
encoder.summary()

4. El decoder

latent_inputs = keras.Input(shape=(latent_dim,))
x = layers.Dense(intermediate_dim, activation="relu")(latent_inputs)
x = layers.Dense(original_dim, activation="sigmoid")(x)
decoder_outputs = layers.Reshape((28, 28, 1))(x)
decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")
decoder.summary()

5. El modelo VAE completo, con su train_step personalizado

Un VAE necesita un bucle de entrenamiento propio porque la pérdida (reconstrucción + KL) no es una pérdida estándar de Keras: se define sobreescribiendo train_step.

class VAE(keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super().__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.total_loss_tracker = keras.metrics.Mean(name="total_loss")
        self.reconstruction_loss_tracker = keras.metrics.Mean(name="reconstruction_loss")
        self.kl_loss_tracker = keras.metrics.Mean(name="kl_loss")

    @property
    def metrics(self):
        return [self.total_loss_tracker, self.reconstruction_loss_tracker, self.kl_loss_tracker]

    def train_step(self, data):
        with tf.GradientTape() as tape:
            z_mean, z_log_var, z = self.encoder(data)
            reconstruction = self.decoder(z)
            reconstruction_loss = tf.reduce_mean(
                tf.reduce_sum(
                    keras.losses.binary_crossentropy(data, reconstruction),
                    axis=(1, 2),
                )
            )
            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
            kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))
            total_loss = reconstruction_loss + kl_loss

        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        self.total_loss_tracker.update_state(total_loss)
        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
        self.kl_loss_tracker.update_state(kl_loss)
        return {
            "loss": self.total_loss_tracker.result(),
            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
            "kl_loss": self.kl_loss_tracker.result(),
        }

6. Entrenamiento

vae = VAE(encoder, decoder)
vae.compile(optimizer=keras.optimizers.Adam())
vae.fit(mnist_digits, epochs=30, batch_size=128)

Vigilar reconstruction_loss y kl_loss por separado en cada época, no solo loss, es lo que permite detectar el fallo de entrenamiento que se explica en la siguiente sección antes de que arruine horas de cómputo.

7. Visualizar el espacio latente

La comprobación de que el espacio latente aprendido es realmente continuo no es teórica: se puede dibujar, porque se fijó latent_dim = 2. Esta función genera una rejilla de puntos (x, y) en el espacio latente y decodifica cada uno en una imagen:

def plot_latent_space(vae, n=30, digit_size=28, scale=1.0, figsize=15):
    figure = np.zeros((digit_size * n, digit_size * n))
    grid_x = np.linspace(-scale, scale, n)
    grid_y = np.linspace(-scale, scale, n)[::-1]

    for i, yi in enumerate(grid_y):
        for j, xi in enumerate(grid_x):
            z_sample = np.array([[xi, yi]])
            x_decoded = vae.decoder.predict(z_sample, verbose=0)
            digit = x_decoded[0].reshape(digit_size, digit_size)
            figure[
                i * digit_size : (i + 1) * digit_size,
                j * digit_size : (j + 1) * digit_size,
            ] = digit

    plt.figure(figsize=(figsize, figsize))
    plt.imshow(figure, cmap="Greys_r")
    plt.axis("off")
    plt.show()

plot_latent_space(vae)

El resultado es una rejilla de 30x30 dígitos donde cada celda es la interpolación de la celda anterior: un 4 se transforma suavemente en un 9, un 0 en un 6, sin saltos abruptos. Esa transición continua es la prueba práctica de que la divergencia KL cumplió su función: el espacio latente no es un conjunto de puntos memorizados, es una superficie donde moverse un poco produce un cambio proporcional en la salida. Es, con diferencia, la pieza más reveladora de todo el ejercicio, más que cualquier métrica de reconstrucción aislada.

Qué falla realmente al entrenar un VAE: el colapso del posterior

El fallo de entrenamiento más citado en la literatura sobre VAE es el colapso del posterior (o KL vanishing): la divergencia KL se acerca a cero, lo que significa que q(z|x) converge al prior p(z) sin importar cuál sea la entrada. El decoder aprende a ignorar z por completo y reconstruye apoyándose en su propia capacidad, sin usar la información latente. El síntoma es un espacio latente no informativo: cualquier punto que muestrees produce salidas casi indistinguibles entre sí, muy distinto de la rejilla continua y variada del apartado anterior.

Este problema se documentó y se resolvió por primera vez con claridad fuera del dominio de imagen: en VAEs para texto, con un decoder autorregresivo (una LSTM) que es tan potente por sí solo que puede reconstruir la secuencia sin apoyarse en z. Bowman et al., Generating Sentences from a Continuous Space, identificaron el problema y propusieron la mitigación que se sigue usando hoy como punto de partida: el KL annealing. En lugar de aplicar el peso completo del término KL desde el primer paso de entrenamiento, se empieza con ese peso en cero (el modelo se comporta como un autocodificador determinista mientras aprende a reconstruir) y se sube gradualmente durante las primeras épocas, dando tiempo a que el encoder aprenda a usar z antes de que la regularización empiece a presionar.

El diagnóstico práctico no requiere herramientas adicionales a las que ya aparecen en el bucle de entrenamiento de más arriba: basta con separar reconstruction_loss y kl_loss en vez de mirar solo loss. Si la pérdida total mejora pero kl_loss se queda pegada cerca de cero mientras reconstruction_loss sigue bajando, hay colapso del posterior en curso, no un modelo que está aprendiendo bien. Una segunda señal complementaria es contar las unidades activas del espacio latente: cuántas dimensiones tienen varianza no trivial entre distintos ejemplos. Un espacio de 20 dimensiones con solo 3 activas es colapso parcial, aunque las métricas de reconstrucción parezcan razonables.

Dos técnicas adicionales, además del KL annealing, aparecen con frecuencia junto a este problema, y conviene tener claro qué hace cada una:

• Beta-VAE pondera el término KL con un factor beta en lugar de dejarlo en 1. Un beta por debajo de 1 relaja la presión de regularización durante el entrenamiento, en la misma dirección que el annealing. Pero beta por encima de 1 hace justo lo contrario de mitigar el colapso: aumenta el peso relativo de la divergencia KL, lo que empuja con más fuerza a q(z|x) hacia el prior y puede agravar el colapso en vez de evitarlo. beta > 1 se usa para otro objetivo distinto, forzar representaciones más disentangled, no para resolver KL vanishing.
• Free bits no fuerza una divergencia KL mínima: deja de penalizarla por debajo de un umbral. Se aplica un margen por dimensión latente de forma que, mientras la KL de esa dimensión esté por debajo del umbral, no contribuye al gradiente. Por encima del umbral, sí. El efecto neto es que el encoder tiene libertad para mantener algo de información en z sin que la regularización la aplaste, pero nada en la formulación obliga a que la KL alcance ese mínimo.

Cuándo elegir un VAE frente a difusión o una GAN

Una vez que el VAE funciona, queda una pregunta distinta y práctica: con los modelos de difusión dominando la síntesis fotorrealista y las GAN todavía compitiendo en dominios estrechos, ¿en qué caso concreto sigue teniendo sentido un VAE en lugar de otra arquitectura? La respuesta depende de cuatro ejes:

• Coste y latencia de muestreo: un VAE genera en un único paso hacia adelante; un modelo de difusión necesita decenas o cientos de pasos de denoising iterativo.
• Estructura del espacio latente: si necesitas interpolar, hacer aritmética de atributos o usar el espacio latente como representación para otra tarea (detección de anomalías, clustering, búsqueda por similitud), importa que ese espacio esté explícitamente regularizado hacia una distribución conocida, como se vio en la rejilla de dígitos.
• Estabilidad de entrenamiento: un VAE optimiza una única pérdida bien definida (la ELBO). Una GAN entrena dos redes en un juego adversarial que puede colapsar sin previo aviso.
• Tolerancia a la borrosidad: si el objetivo es fidelidad perceptual máxima, un VAE puro pierde frente a difusión; si el objetivo es una representación útil o una reconstrucción aproximada rápida, la borrosidad es un coste aceptable.

Esta última fila no es un caso teórico: los Latent Diffusion Models descritos en Rombach et al., la base de Stable Diffusion, usan exactamente un VAE para comprimir a un espacio latente compacto y aplican difusión dentro de ese espacio en lugar de sobre los píxeles crudos. En ese sentido, la pregunta rara vez es "VAE o difusión": en producción, casi siempre es "VAE y difusión juntos, o difusión sola sobre píxeles". Para comparar cualquiera de estas familias de forma objetiva, más allá de mirar imágenes, la métrica estándar es FID (Fréchet Inception Distance), introducida en Heusel et al., que compara la distribución de activaciones de una red Inception entre datos reales y generados y penaliza tanto la falta de fidelidad como la falta de diversidad.

Señales de que tu VAE está aprendiendo bien

Antes de considerar terminado un entrenamiento, conviene revisar tres cosas concretas sobre el propio modelo que acabas de construir, en este orden:

• ¿kl_loss se mantiene en un rango razonable durante todo el entrenamiento, o cae a un valor cercano a cero mientras reconstruction_loss sigue bajando? Lo segundo es colapso del posterior, no una buena señal.
• ¿Cuántas dimensiones del espacio latente tienen varianza no trivial? Con latent_dim = 2 es fácil de comprobar a ojo con plot_latent_space; con dimensiones más altas, hace falta medir la varianza por dimensión explícitamente.
• ¿La rejilla de interpolación muestra transiciones suaves entre clases, o regiones enteras que producen la misma salida borrosa sin importar el punto exacto? La segunda es la versión visual del mismo problema.

Estas tres comprobaciones cuestan minutos y evitan la situación más frustrante de entrenar un VAE: pasar horas de cómputo para terminar con un modelo que technically converge pero que, en la práctica, no aprendió nada sobre la estructura de los datos.