Auto-Encoding Variational Bayes

Guía de lectura

Problema y motivación

Los modelos probabilísticos dirigidos con variables latentes continuas $\mathbf{z}$ son útiles para el modelado generativo, el aprendizaje de representaciones no supervisado y la inferencia marginal. Sin embargo, su aprendizaje mediante máxima verosimilitud o métodos bayesianos enfrenta graves obstáculos computacionales cuando:

1. Intratabilidad del posterior: Para relaciones no lineales (por ejemplo, redes neuronales profundas), la integral de la verosimilitud marginal $p_{\mathbf{\theta }}(\mathbf{x})=\int p_{\mathbf{\theta }}(\mathbf{z})p_{\mathbf{\theta }}(\mathbf{x}|\mathbf{z})d\mathbf{z}$ y el posterior $p_{\mathbf{\theta }}(\mathbf{z}|\mathbf{x})$ carecen de soluciones analíticas en forma cerrada, haciendo inaplicables el algoritmo clásico de Expectation-Maximization (EM) y el método variacional de campo medio (mean-field variational Bayes).
2. Grandes conjuntos de datos: Los enfoques tradicionales de muestreo, como Monte Carlo EM, requieren un muestreo MCMC lento e iterativo por cada punto de datos a lo largo de todo el conjunto de datos, lo que hace prohibitivamente costosa la optimización por mini-lotes a gran escala.

El estimador de gradiente ingenuo estándar basado en la función de puntuación (score-function) ${\nabla }_{\mathbf{\varphi }}{\mathbb{E}}_{q_{\mathbf{\varphi }}(\mathbf{z})}[f(\mathbf{z})]$ exhibe una varianza excesivamente alta, lo que impide una optimización fiable con ascenso de gradiente estocástico.

Ideas clave

• Reparameterization Trick (Truco de reparametrización): En lugar de muestrear $\mathbf{z}$ directamente de una distribución parametrizada intratable $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$, la variable se reescribe como una transformación determinista de una variable de ruido independiente $\mathbf{ϵ}\sim p(\mathbf{ϵ})$: $\mathbf{z}=g_{\mathbf{\varphi }}(\mathbf{ϵ},\mathbf{x})$. Esto permite la diferenciación con respecto a $\mathbf{\varphi }$ a través de la retropropagación estándar con una reducción drástica de la varianza.
• Amortized Inference (Inferencia amortizada o modelo de reconocimiento): En lugar de optimizar parámetros variacionales por cada punto de datos individual, una red neuronal codificadora $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$ predice los parámetros variacionales (como la media y la varianza de una gaussiana latente) condicionados a $\mathbf{x}$.
• Auto-Encoding Variational Bayes (AEVB): Optimiza conjuntamente los parámetros generativos $\mathbf{\theta }$ (decodificador probabilístico) y los parámetros variacionales $\mathbf{\varphi }$ (codificador probabilístico) mediante ascenso de gradiente estocástico sobre el límite inferior variacional reparametrizado.

Método

El objetivo variacional maximiza el Evidence Lower Bound (ELBO) $\mathcal{L}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})$: $$\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)})\ge \mathcal{L}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})=-D_{KL}(q_{\mathbf{\varphi }}(\mathbf{z}|{\mathbf{x}}^{(i)})\parallel p_{\mathbf{\theta }}(\mathbf{z}))+{\mathbb{E}}_{q_{\mathbf{\varphi }}(\mathbf{z}|{\mathbf{x}}^{(i)})}\left[\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)}|\mathbf{z})\right]$$

Aplicando el truco de reparametrización ${\mathbf{z}}^{(i,l)}=g_{\mathbf{\varphi }}({\mathbf{ϵ}}^{(i,l)},{\mathbf{x}}^{(i)})$ con ${\mathbf{ϵ}}^{(i,l)}\sim p(\mathbf{ϵ})$, el estimador genérico SGVB (Estimador A) es: $${\tilde{\mathcal{L}}}^A(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})=\frac{1}{L}\sum _{l=1}^L\left[\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)},{\mathbf{z}}^{(i,l)})-\log q_{\mathbf{\varphi }}({\mathbf{z}}^{(i,l)}|{\mathbf{x}}^{(i)})\right]$$

Cuando el prior $p_{\mathbf{\theta }}(\mathbf{z})=\mathcal{N}(0,\mathbf{I})$ y el posterior aproximado $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})=\mathcal{N}(\mathbf{\mu },{\mathbf{\sigma }}^2\mathbf{I})$ son gaussianos, la divergencia KL se puede calcular analíticamente, dando lugar al Estimador B: $$\mathcal{L}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})\simeq \frac{1}{2}\sum _{j=1}^J\left(1+\log (({\sigma }_j^{(i)}{)}^2)-({\mu }_j^{(i)}{)}^2-({\sigma }_j^{(i)}{)}^2\right)+\frac{1}{L}\sum _{l=1}^L\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)}|{\mathbf{z}}^{(i,l)})$$ donde ${\mathbf{z}}^{(i,l)}={\mathbf{\mu }}^{(i)}+{\mathbf{\sigma }}^{(i)}\odot {\mathbf{ϵ}}^{(l)}$ y ${\mathbf{ϵ}}^{(l)}\sim \mathcal{N}(0,\mathbf{I})$.

Para mini-lotes de tamaño $M$, el estimador completo del límite inferior es: $${\tilde{\mathcal{L}}}^M(\mathbf{\theta },\mathbf{\varphi };{\mathbf{X}}^M)=\frac{N}{M}\sum _{i=1}^M\tilde{\mathcal{L}}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})$$ En la práctica, fijar $L=1$ muestra por punto de datos es suficiente cuando se combina con tamaños de mini-lote como $M=100$.

Resultados

• Configuración de los modelos: Los modelos generativos se evaluaron en MNIST (decodificador MLP Bernoulli, 500 unidades ocultas) y Frey Face (decodificador MLP gaussiano, 200 unidades ocultas), optimizados mediante Adagrad con tasas de aprendizaje en $\{0.01,0.02,0.1\}$ y decaimiento de pesos (weight decay).
• Optimización del límite inferior: En diversas dimensionalidades latentes ($N_{\mathbf{z}}\in \{2,5,10,20,400\}$ en Frey Face y MNIST), AEVB convergió significativamente más rápido y alcanzó un límite inferior más alto que el algoritmo wake-sleep.
• Resistencia al sobreajuste: Aumentar la dimensionalidad de $\mathbf{z}$ no provocó sobreajuste debido al efecto regularizador del término analítico de la divergencia KL respecto al prior.
• Estimación de la verosimilitud marginal: En experimentos con un espacio latente tridimensional en MNIST utilizando un estimador de verosimilitud marginal basado en HMC, AEVB convergió más rápido que wake-sleep y superó ampliamente la viabilidad de Monte Carlo EM (que resultaba demasiado lento para ejecutarse en el conjunto de datos MNIST completo).
• Tiempo de ejecución: El entrenamiento tomó aproximadamente entre 20 y 40 minutos por cada millón de muestras en una CPU Intel Xeon (~40 GFLOPS).

Limitaciones y preguntas abiertas

• Variables latentes continuas: El truco de reparametrización depende directamente de la diferenciabilidad con respecto a los parámetros, lo que hace que no sea trivial aplicarlo de forma directa a distribuciones de variables latentes discretas.
• Brecha de aproximación (Approximation Gap): Asumir una distribución gaussiana diagonal para $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$ limita la capacidad expresiva cuando el posterior real es multimodal o presenta un acoplamiento fuerte.
• Extensión a Variational Bayes completo: La evaluación empírica se centró en la estimación de parámetros MAP/ML para $\mathbf{\theta }$; la inferencia variacional sobre los parámetros globales $\mathbf{\theta }$ se formuló en el apéndice pero se dejó pendiente para una futura validación empírica.