Auto-Encoding Variational Bayes

Guia de leitura

Problema e motivação

Modelos probabilísticos direcionados com variáveis latentes contínuas $\mathbf{z}$ são úteis para modelagem generativa, aprendizado de representações não supervisionado e inferência marginal. No entanto, aprendê-los via máxima verossimilhança ou métodos bayesianos enfrenta severos obstáculos computacionais quando:

1. Intratabilidade da posteriori: Para relações não lineares (por exemplo, redes neurais profundas), a integral da verossimilhança marginal $p_{\mathbf{\theta }}(\mathbf{x})=\int p_{\mathbf{\theta }}(\mathbf{z})p_{\mathbf{\theta }}(\mathbf{x}|\mathbf{z})d\mathbf{z}$ e a posteriori $p_{\mathbf{\theta }}(\mathbf{z}|\mathbf{x})$ não possuem soluções analíticas em forma fechada, tornando o Expectation-Maximization (EM) padrão e o variational Bayes de campo médio inaplicáveis.
2. Grandes conjuntos de dados: Abordagens tradicionais de amostragem, como o Monte Carlo EM, exigem uma amostragem MCMC lenta e iterativa por ponto de dado em todo o conjunto de dados, tornando a otimização em mini-lotes em grande escala proibitivamente cara.

O estimador de gradiente ingênuo padrão por função de escore ${\nabla }_{\mathbf{\varphi }}{\mathbb{E}}_{q_{\mathbf{\varphi }}(\mathbf{z})}[f(\mathbf{z})]$ exibe variância excessivamente alta, impedindo uma otimização confiável com subida de gradiente estocástica.

Ideias centrais

• Truque da Reparametrização (Reparameterization Trick): Em vez de amostrar $\mathbf{z}$ diretamente de uma distribuição parametrizada intratável $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$, a variável é reescrita como uma transformação determinística de uma variável de ruído independente $\mathbf{ϵ}\sim p(\mathbf{ϵ})$: $\mathbf{z}=g_{\mathbf{\varphi }}(\mathbf{ϵ},\mathbf{x})$. Isso permite a diferenciação em relação a $\mathbf{\varphi }$ por meio de retropropagação padrão com variância drasticamente reduzida.
• Inferência Amortizada (Modelo de Reconhecimento): Em vez de otimizar parâmetros variacionais por ponto de dado, uma rede neural codificadora $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$ prediz os parâmetros variacionais (como a média e a variância de uma gaussiana latente) condicionados a $\mathbf{x}$.
• Auto-Encoding Variational Bayes (AEVB): Otimiza conjuntamente os parâmetros generativos $\mathbf{\theta }$ (decodificador probabilístico) e os parâmetros variacionais $\mathbf{\varphi }$ (codificador probabilístico) usando subida de gradiente estocástica no limite inferior variacional reparametrizado.

Método

O objetivo variacional maximiza o Limite Inferior da Evidência (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 o truque da reparametrização ${\mathbf{z}}^{(i,l)}=g_{\mathbf{\varphi }}({\mathbf{ϵ}}^{(i,l)},{\mathbf{x}}^{(i)})$ com ${\mathbf{ϵ}}^{(i,l)}\sim p(\mathbf{ϵ})$, o estimador SGVB genérico (Estimador A) é: $${\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]$$

Quando a priori $p_{\mathbf{\theta }}(\mathbf{z})=\mathcal{N}(0,\mathbf{I})$ e a posteriori aproximada $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})=\mathcal{N}(\mathbf{\mu },{\mathbf{\sigma }}^2\mathbf{I})$ são gaussianas, a divergência KL pode ser calculada analiticamente, resultando no 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)})$$ onde ${\mathbf{z}}^{(i,l)}={\mathbf{\mu }}^{(i)}+{\mathbf{\sigma }}^{(i)}\odot {\mathbf{ϵ}}^{(l)}$ e ${\mathbf{ϵ}}^{(l)}\sim \mathcal{N}(0,\mathbf{I})$.

Para mini-lotes de tamanho $M$, o estimador completo do limite inferior é: $${\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)})$$ Na prática, definir $L=1$ amostra por ponto de dado é suficiente quando combinado com tamanhos de mini-lote como $M=100$.

Resultados

• Configurações do Modelo: Os modelos generativos foram avaliados no MNIST (decodificador MLP Bernoulli, 500 unidades ocultas) e Frey Face (decodificador MLP Gaussiano, 200 unidades ocultas), otimizados via Adagrad com taxas de aprendizado em $\{0.01,0.02,0.1\}$ e decaimento de pesos (weight decay).
• Otimização do Limite Inferior: Em várias dimensionalidades latentes ($N_{\mathbf{z}}\in \{2,5,10,20,400\}$ no Frey Face e MNIST), o AEVB convergiu significativamente mais rápido e atingiu um limite inferior mais alto do que o algoritmo wake-sleep.
• Resistência ao Sobreajuste: Aumentar a dimensionalidade de $\mathbf{z}$ não causou sobreajuste devido ao efeito regularizador do termo analítico da divergência KL em direção à priori.
• Estimação da Verossimilhança Marginal: Em experimentos com espaço latente tridimensional no MNIST usando um estimador de verossimilhança marginal baseado em HMC, o AEVB convergiu mais rápido que o wake-sleep e superou amplamente a viabilidade do Monte Carlo EM (que era lento demais para ser executado em todo o conjunto de dados MNIST).
• Tempo de Execução: O treinamento levou aproximadamente de 20 a 40 minutos por milhão de amostras em uma CPU Intel Xeon (~40 GFLOPS).

Limitações e questões em aberto

• Variáveis Latentes Contínuas: O truque da reparametrização depende diretamente da diferenciabilidade em relação aos parâmetros, tornando sua aplicação direta a distribuições de variáveis latentes discretas algo não trivial.
• Gap de Aproximação: Assumir uma distribuição gaussiana diagonal para $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$ limita a capacidade expressiva quando a verdadeira posteriori é multimodal ou fortemente acoplada.
• Extensão para Variational Bayes Completo: A avaliação empírica focou na estimação de parâmetros por MAP/ML para $\mathbf{\theta }$; a inferência variacional sobre os parâmetros globais $\mathbf{\theta }$ foi formulada no apêndice, mas deixada para validação empírica futura.