17. VAE

O VAE é o capítulo que as pessoas mais leem errado. Julgado como gerador de imagens ele perdeu, de forma decisiva e suas amostras borradas viraram piada. Julgado pelo uso em produção, é um dos modelos generativos mais empregados que existem: toda imagem de difusão latente que você já viu foi decodificada por um deles.

As duas afirmações são verdadeiras e a razão de as duas serem verdadeiras é o ponto deste capítulo. Um VAE não é bem um gerador. É um compressor com um espaço latente moldado e assim que você o enxerga desse jeito o borrão deixa de ser defeito e vira uma escolha de projeto que outra coisa corrige adiante.

Comece por um autoencoder e veja o que falta

Um autoencoder são duas redes prensadas contra um gargalo: um encoder \(E\) que leva a entrada a um código de baixa dimensão e um decoder \(D\) que traz de volta.

Nada obriga os códigos a se organizarem de forma sensata. A reconstrução só se importa que cada entrada vá para algum código que o decodificador consiga inverter, então os códigos se espalham para onde o treino os empurrou, deixando buracos entre eles. Escolha um ponto aleatório nesse espaço e o decodificador produz lixo, porque ele nunca foi treinado por perto.

Duas variantes que vale nomear

Autoencoders convolucionais usam CNNs nas duas metades. Para imagens isso não é detalhe — é o prior certo (capítulo 10) e é o que todo autoencoder latente de produção ainda usa.

Autoencoders com remoção de ruído e mascarados corrompem a entrada e pedem a versão limpa de volta. Essa mudança de objetivo é o que transformou autoencoders em ferramenta séria de aprendizado de representações — o MAE é essa ideia em escala de ViT e modelos de difusão são, de longe, a mesma ideia com um contínuo de níveis de ruído.

O VAE: codifique uma distribuição, não um ponto

A correção é fazer o encoder produzir uma distribuição sobre códigos em vez de um único código e então empurrar essas distribuições na direção de um prior do qual você sabe amostrar.

O ELBO, um passo de cada vez

A perda do VAE não é uma receita que alguém ajustou até funcionar. Ela sai de uma pergunta e de uma identidade. Cada termo dela tem uma função que dá para nomear. São cinco passos.

1. O que gostaríamos de maximizar

Um modelo generativo é uma história de como os dados são feitos: sorteie um código \(z \sim p(z)\) e depois sorteie uma imagem do decoder, \(x \sim p_\theta(x \mid z)\). A melhor história é a que torna mais prováveis os dados que de fato temos — máxima verossimilhança, o mesmo princípio por trás de toda perda deste curso. Para uma imagem, a quantidade é

Esse número se chama evidência (evidence): o quanto o modelo acredita em \(x\). É o E de ELBO.

O problema é a integral sobre todos os códigos possíveis. Dá para tentar estimá-la sorteando códigos do prior e tirando a média de \(p_\theta(x \mid z)\), só que quase todo \(z\) sorteado do prior não tem nada a ver com este \(x\) — um código aleatório decodifica em outro dígito — e a média acaba decidida por um punhado de amostras raras e sortudas. Os códigos que importam para \(x\) são os que a posterior \(p(z \mid x) = p_\theta(x \mid z)\,p(z)\,/\,p_\theta(x)\) aponta. Só que calcular a posterior exige \(p_\theta(x)\), justamente o número que procurávamos. O problema é circular.

2. Chute a posterior e depois meça o chute

Quebre o círculo com uma aproximação. O encoder produz um chute gaussiano para a posterior,

("variacional" quer dizer apenas escolher o melhor membro de uma família de distribuições). Agora escreva a evidência usando \(q\). Cada linha é uma igualdade, não uma aproximação:

3. Separe nos dois termos que você vai programar

Escreva \(p_\theta(x, z) = p_\theta(x \mid z)\,p(z)\) dentro do ELBO e ele se separa em duas esperanças que conseguimos calcular:

4. Um ELBO que dá para calcular na mão

Pegue um modelo pequeno o bastante para ser resolvido exatamente: um código 1-D, o prior \(p(z) = \mathcal{N}(0, 1)\), um decoder que multiplica por 2 e soma ruído unitário (\(p(x \mid z) = \mathcal{N}(2z, 1)\)) e uma observação, \(x = 3\). Como tudo é gaussiano, os dois números que são intratáveis num VAE de verdade estão disponíveis aqui:

• Evidência. \(x = 2z + \text{ruído}\) é gaussiano com variância \(2^2 + 1 = 5\), então \(\log p(3) = -\tfrac12 \log(2\pi \cdot 5) - \tfrac{9}{10} = -2.6237\).
• Posterior. \(p(z \mid x = 3) = \mathcal{N}\big(\tfrac{2 \cdot 3}{5},\ \tfrac15\big) = \mathcal{N}(1.2,\ 0.2)\).

Para um chute \(q = \mathcal{N}(m, v)\) os dois termos do ELBO têm forma fechada:

Repare no \(4v\): quanto mais largo o chute, mais o código sorteado erra e mais a reconstrução paga. É o cabo de guerra num símbolo só — a reconstrução quer \(v \to 0\) e a KL quer \(v \to 1\).

Uma linha por extenso, para o chute \(q_C = \mathcal{N}(1,\ 0.5^2)\):

e quatro chutes lado a lado:

Leia as linhas como os três personagens deste capítulo:

• O prior paga KL zero e a reconstrução dele é péssima: o chute não diz nada sobre \(x\). É o colapso do posterior em miniatura.
• O ponto \(q_B\) tem a melhor reconstrução de todas — ele fica em \(z = x/2 = 1.5\), que decodifica exatamente em \(x\) — e paga 2.93 nats de KL por ser tão estreito. É o autoencoder e o ELBO dele é pior que o de um chute modesto como \(q_C\).
• A posterior não tem nem a melhor reconstrução nem a menor KL. Ela tem o melhor ELBO e esse ELBO é exatamente \(\log p(x)\): a folga fechou.

O ELBO não premia a melhor reconstrução. Ele premia a melhor troca.

5. Do ELBO à perda que você programa

Troque o sinal — otimizadores minimizam — e faça três substituições:

1. A esperança vira uma amostra. \(\mathbb{E}_q[\cdot]\) é estimada com um único \(z = \mu + \sigma\epsilon\) por imagem — o truque da reparametrização, logo abaixo. O passo 2 do lab mostra por que uma amostra basta.
2. \(-\log p_\theta(x \mid z)\) vira uma perda conhecida, conforme o que se assume sobre os pixels. Um decoder gaussiano \(\mathcal{N}(\hat{x}, \sigma_x^2 I)\) dá \(\frac{1}{2\sigma_x^2}\|x - \hat{x}\|^2 + \text{const}\) — o erro quadrático. Um decoder de Bernoulli (pixels em \([0,1]\), saída sigmoide) dá a entropia cruzada binária. A fórmula abaixo descarta o \(\frac{1}{2\sigma_x^2}\); escolher outro \(\sigma_x\) é o mesmo que reescalar o \(\beta\).
3. A KL é escrita em forma fechada. Para uma dimensão, \(\text{KL}\big(\mathcal{N}(\mu, \sigma^2)\,\|\,\mathcal{N}(0,1)\big) = \mathbb{E}_q\big[\log q(z) - \log p(z)\big] = \mathbb{E}_q\Big[-\tfrac12\log\sigma^2 - \tfrac{(z-\mu)^2}{2\sigma^2} + \tfrac{z^2}{2}\Big] = \tfrac12\big(\mu^2 + \sigma^2 - 1 - \log\sigma^2\big)\), usando \(\mathbb{E}_q[(z-\mu)^2] = \sigma^2\) e \(\mathbb{E}_q[z^2] = \mu^2 + \sigma^2\). As duas distribuições são diagonais, então as dimensões simplesmente se somam.

O resultado é a perda que toda implementação de VAE calcula:

Com \(\beta = 1\) ela é exatamente o ELBO negativo. A KL só tem forma fechada porque as duas distribuições são gaussianas e essa é a única razão pela qual tudo isso é prático. E o \(\beta\) — o peso entre os dois termos — é onde mora todo o comportamento.

O β é a história inteira

O painel abaixo treina um VAE de verdade no seu navegador: duas camadas ocultas, o truque da reparametrização, Adam. Mova o \(\beta\), aperte retreinar e observe três coisas ao mesmo tempo — quão bem ele reconstrói, que forma o latente assume e a única pergunta que importa para geração: uma amostra de \(\mathcal{N}(0, I)\) decodifica em algo real?

Três regimes e todo VAE que você treinar estará em algum ponto entre eles:

• \(\beta \to 0\) — um autoencoder comum. A reconstrução é quase perfeita e cerca de 90% das amostras do prior decodificam em nada. É a falha que o termo KL existe para evitar, tornada visível.
• \(\beta\) numa faixa boa e estreita — o latente casa com o prior o bastante para amostrar e ainda carrega informação suficiente para reconstruir.
• \(\beta\) grande demais — colapso do posterior. A forma mais barata de satisfazer uma penalidade KL grande é o encoder ignorar a entrada e devolver \(\mathcal{N}(0,I)\) para tudo. A KL vai a zero, o decodificador não tem em que se condicionar e ele emite a média do conjunto. A perda parece estar caindo.

Você já encontrou os três no exemplo resolvido. O ponto \(q_B\) é aonde \(\beta \to 0\) leva: só reconstrução, nenhuma consideração pelo prior. O próprio prior é aonde \(\beta \to \infty\) leva: KL zero e um código que não diz nada. A posterior verdadeira é a resposta para \(\beta = 1\). O passo 4 do lab resolve o melhor \(q\) para cada \(\beta\) em forma fechada e o vê deslizar de um extremo ao outro.

O truque da reparametrização

Amostrar não é diferenciável, então um VAE ingênuo não treina por retropropagação, ponto. A correção é tirar a aleatoriedade do caminho por onde o gradiente passa:

Agora \(z\) é função determinística de \(\mu\), \(\sigma\) e de uma entrada \(\epsilon\) que não carrega parâmetro nenhum. Gradientes fluem para \(\mu\) e \(\sigma\); o ruído entra como dado. Essa única linha é por que o artigo do VAE⁸ é de 2013 e não de uma década antes e o mesmo truque aparece por todo o aprendizado de máquina onde um nó estocástico fica dentro de um modelo diferenciável.

Concretamente, o gradiente do termo de reconstrução agora chega às duas saídas do encoder pela regra da cadeia:

A segunda carrega o cabo de guerra: qualquer que seja o sinal de \(\epsilon\), um ruído que piorou a reconstrução empurra \(\sigma\) para baixo, enquanto o termo KL o puxa de volta para 1. O passo 3 do lab roda essa atualização (uma amostra por vez) e chega à posterior verdadeira.

Uma passagem de VAE, número por número

Tudo até aqui — as duas cabeças, o truque, os dois termos da perda — numa rede pequena o bastante para acompanhar com lápis: duas entradas, um encoder de dois neurônios, um latente de uma dimensão, um decoder de dois neurônios e duas saídas. Percorra a passagem para a frente e depois a de volta. As cores são as do capítulo de MLP: verde para a passagem direta, laranja para gradientes e roxo para pesos.

Três coisas para reparar no caminho:

• Os passos 2, 3, 5 e 6 são um MLP comum. Um VAE são dois MLPs com um número aleatório no meio; toda regra de retropropagação que você já conhece vale sem mudança.
• O passo 4 é a única operação nova e o passo 10 é a derivada dele: \(\partial z/\partial \mu = 1\) e \(\partial z/\partial s = \tfrac12\sigma\epsilon\). Sem o truque, o passo 10 não teria pelo que multiplicar e o encoder não receberia gradiente nenhum da reconstrução.
• O passo 10 é onde mora o cabo de guerra. Arraste o \(\epsilon\) e observe a parte da reconstrução em \(\partial L/\partial s\): ela troca de sinal de um sorteio para outro, mas a média sobre \(\epsilon\), mostrada na legenda, é positiva — o ruído custa reconstrução, então empurra \(\sigma\) para baixo. A parte da KL é a mesma para todo \(\epsilon\) e sempre empurra \(\sigma\) para 1.

Por que as amostras são borradas — e por que isso não importa

Volte ao capítulo 16. O ELBO é um limite da verossimilhança, então um VAE maximiza verossimilhança, então ele otimiza a KL para frente, que é cobridora de modos. Quando o modelo não consegue representar os dados exatamente, ele hedgeia — colocando probabilidade entre os modos em vez de escolher um.

Some a versão mecânica: uma perda de reconstrução \(\ell_2\) é uma verossimilhança gaussiana e a predição ótima sob erro quadrático é a média condicional. Se várias imagens nítidas são compatíveis com um código, a saída que minimiza a perda é a média delas. Médias de imagens são borradas. Não há bug aqui; o objetivo recebeu exatamente o que pediu.

Então o borrão se corrige mudando o objetivo, que é precisamente o que um autoencoder latente de produção faz:

O termo perceptual compara features profundas em vez de pixels, então uma textura deslocada em um pixel não é punida. O termo adversarial — um discriminador de GAN sobre a saída do decodificador (capítulo 18) — pune exatamente a suavidade que o \(\ell_2\) premia. Todo autoencoder latente de todo modelo de difusão é treinado assim.

Para que serve um VAE, de fato

Aqui está a história de produção e não é a que o capítulo de livro-texto sugere.

Uma imagem RGB 512×512 são 786.432 números. Rodar difusão ali é inviável. A difusão latente⁹ roda num espaço comprimido — e um VAE convolucional é o que comprime:

O modelo generativo nunca vê um pixel. Ele trabalha inteiramente no espaço latente do VAE e o decodificador transforma o resultado em imagem só no final. Três consequências que vale carregar:

1. O VAE limita sua qualidade. O que o encoder não representa, o modelo de difusão não gera. Texto pequeno e textura fina são perdas conhecidas — e a mudança de 4 para 16 canais latentes no SD3 e no FLUX mirou exatamente nisso.
2. O β dele é minúsculo. Estes são compressores: não se amostra deles diretamente, então o termo KL só precisa manter a escala do latente limitada e razoavelmente regular. Valores práticos ficam na casa de \(10^{-6}\).
3. Ele fica congelado. Treine o autoencoder uma vez e o mantenha fixo enquanto treina o modelo de difusão. É por isso que um ecossistema inteiro de LoRAs e fine-tunes compartilha um mesmo espaço latente.

O outro ramo: latentes discretos

Em vez de um latente gaussiano, o VQ-VAE¹⁰ encaixa cada vetor latente na entrada mais próxima de um codebook aprendido. O latente vira uma grade de inteiros.

É um projeto diferente com propósito diferente e a razão de conhecê-lo é que ele é o que torna o capítulo 23 possível: uma vez que uma imagem é uma grade de inteiros, ela é uma sequência de tokens e dá para rodar um Transformer sobre ela exatamente como sobre texto.

Pontos principais

1. Um autoencoder comprime; ele não gera. O espaço latente dele tem buracos e você não os enxerga.
2. Um VAE codifica uma distribuição por entrada e puxa essas distribuições na direção de um prior. É isso que torna o espaço amostrável.
3. \(\log p(x) = \text{ELBO} + \text{KL}(q \,\|\, p(z \mid x))\). O ELBO é um limite inferior da evidência, a folga é o quanto o encoder erra e subi-lo treina encoder e decoder de uma vez. Por extenso, ele é reconstrução menos KL-até-o-prior: os dois termos são necessários e puxam em direções opostas. Ele premia a melhor troca, não a melhor reconstrução.
4. O β controla tudo. Pequeno demais e é um autoencoder; grande demais e o posterior colapsa. Registre os dois termos separadamente.
5. O truque da reparametrização — \(z = \mu + \sigma\epsilon\) — é o que torna a coisa toda diferenciável.
6. O borrão é o objetivo, não um bug: a KL para frente cobre modos e o \(\ell_2\) devolve a média condicional. Autoencoders de produção corrigem isso com perdas perceptual e adversarial.
7. Em 2026, um VAE é quase sempre o compressor congelado embaixo de um modelo de difusão, a ~48× de compressão e com β minúsculo. Ele limita a qualidade de tudo que é construído em cima.
8. O VQ-VAE é o ramo discreto e é o que torna a geração autorregressiva de imagens possível.

Lab: o ELBO, medido

Quatro experimentos curtos sobre o modelo do exemplo resolvido: \(p(z) = \mathcal{N}(0, 1)\), \(p(x \mid z) = \mathcal{N}(2z, 1)\), uma observação \(x = 3\). Ele é pequeno o bastante para que a evidência \(\log p(x) = -2.6237\) e a posterior \(p(z \mid x) = \mathcal{N}(1.2,\ 0.2)\) sejam conhecidas exatamente, o que nunca acontece num VAE de verdade — e esse é o ponto. Cada afirmação deste capítulo vira um número que dá para conferir. Tudo roda só com numpy: clone os scripts e mude as constantes.

Passo 1 — a folga é o erro da posterior

Sorteie 10 000 chutes aleatórios \(q = \mathcal{N}(m, v)\) e calcule, para cada um, a folga \(\log p(x) - \text{ELBO}\) e, separadamente, \(\text{KL}(q \,\|\, p(z \mid x))\) a partir da posterior exata.

"""Lab 1 — the gap between log p(x) and the ELBO is exactly KL(q || p(z|x)).

Toy model: p(z) = N(0, 1), p(x | z) = N(w z, s^2), w = 2, s = 1, x = 3.
Everything is Gaussian, so the evidence, the posterior and both KLs are exact.
The script draws 10 000 random q = N(m, v), computes the gap two ways, and
checks that no q ever beats log p(x).

Printed as a markdown table, in identifiers only, so one artifact serves both
the English and the Portuguese page.
"""

import numpy as np

W, S, X = 2.0, 1.0, 3.0
VAR_X = W**2 + S**2
LOG_PX = -0.5 * np.log(2 * np.pi * VAR_X) - X**2 / (2 * VAR_X)
M_POST, V_POST = W * X / VAR_X, S**2 / VAR_X      # p(z | x) = N(1.2, 0.2)


def elbo(m, v):
    rec = -0.5 * np.log(2 * np.pi * S**2) - ((X - W * m) ** 2 + W**2 * v) / (2 * S**2)
    kl = 0.5 * (m**2 + v - 1 - np.log(v))
    return rec - kl


def kl_gauss(m1, v1, m2, v2):
    """KL(N(m1, v1) || N(m2, v2))."""
    return 0.5 * (np.log(v2 / v1) + (v1 + (m1 - m2) ** 2) / v2 - 1)


rng = np.random.default_rng(42)
m = rng.uniform(-2, 3, 10_000)
v = np.exp(rng.uniform(-5, 1, 10_000))
gap = LOG_PX - elbo(m, v)
kl_post = kl_gauss(m, v, M_POST, V_POST)

print("| `q = N(m, v)` | `ELBO` | `log p(x) - ELBO` | `KL(q ‖ p(z∣x))` |")
print("|---|---:|---:|---:|")
for i in range(5):
    print(f"| `N({m[i]:+.2f}, {v[i]:.3f})` | {elbo(m[i], v[i]):.4f} | {gap[i]:.4f} | {kl_post[i]:.4f} |")
print(f"| `N({M_POST:.2f}, {V_POST:.3f})` = `p(z∣x)` | **{elbo(M_POST, V_POST):.4f}** | 0.0000 | 0.0000 |")
print()
print("| `10 000 random q` | |")
print("|---|---:|")
print(f"| `log p(x)` | {LOG_PX:.4f} |")
print(f"| `max ELBO found` | {elbo(m, v).max():.4f} |")
print(f"| `q with ELBO > log p(x)` | **{int((gap < 0).sum())}** |")
print(f"| `max ∣gap - KL(q ‖ p(z∣x))∣` | {np.abs(gap - kl_post).max():.1e} |")

As duas últimas colunas da primeira tabela são o mesmo número e a maior discordância em 10 000 chutes é \(10^{-14}\) — ruído de ponto flutuante. É o passo 2 da derivação conferido numericamente: a folga é a KL até a posterior. Nenhum chute passa de \(\log p(x)\) e o melhor dos 10 000 aleatórios para 0.0003 abaixo dele.

Passo 2 — uma amostra basta

Um VAE nunca calcula \(\mathbb{E}_q[\log p(x \mid z)]\) exatamente. Ele sorteia \(z = \mu + \sigma\epsilon\) e tira a média de \(K\) amostras. Quão ruidoso é isso e de quantas amostras ele precisa? O script repete cada estimativa 2 000 vezes para o chute \(q_C = \mathcal{N}(1,\ 0.5^2)\), para o ELBO e para o gradiente dele em relação a \(m\).

"""Lab 2 — estimating the ELBO and its gradient with K samples of z.

Same toy model (p(z) = N(0,1), p(x|z) = N(2z, 1), x = 3), q = N(1, 0.5^2).
A VAE never computes E_q[log p(x|z)] exactly: it samples z = m + sigma * eps
and averages. Each estimate is repeated 2 000 times to measure its mean and
its spread. The KL term is kept in closed form, as in a real VAE.

Printed as a markdown table, in identifiers only, so one artifact serves both
the English and the Portuguese page.
"""

import numpy as np

W, S, X = 2.0, 1.0, 3.0
M, V = 1.0, 0.25
SIGMA = np.sqrt(V)
KL = 0.5 * (M**2 + V - 1 - np.log(V))
EXACT_REC = -0.5 * np.log(2 * np.pi * S**2) - ((X - W * M) ** 2 + W**2 * V) / (2 * S**2)
EXACT_ELBO = EXACT_REC - KL
EXACT_GRAD = W * (X - W * M) / S**2 - M          # d ELBO / d m


def log_lik(z):
    return -0.5 * np.log(2 * np.pi * S**2) - (X - W * z) ** 2 / (2 * S**2)


rng = np.random.default_rng(42)
REPEATS = 2_000

print("| `K` (samples of z) | `ELBO estimate: mean` | `std` | `dELBO/dm estimate: mean` | `std` |")
print("|---:|---:|---:|---:|---:|")
for k in (1, 10, 100, 1000):
    eps = rng.standard_normal((REPEATS, k))
    z = M + SIGMA * eps                               # the reparameterization trick
    elbo = log_lik(z).mean(1) - KL
    grad = (W * (X - W * z) / S**2).mean(1) - M       # dz/dm = 1, KL part = -m
    print(f"| {k} | {elbo.mean():.4f} | {elbo.std():.4f} | {grad.mean():.4f} | {grad.std():.4f} |")
print(f"| `exact` | **{EXACT_ELBO:.4f}** | 0 | **{EXACT_GRAD:.4f}** | 0 |")

A estimativa é não viesada em todo \(K\): as médias caem sobre os valores exatos, \(-2.7371\) e \(1.0\), a menos do ruído de 2 000 repetições. A dispersão cai como \(1/\sqrt{K}\) — 1.29, 0.38, 0.12, 0.04 — ou seja, cem vezes mais amostras compram dez vezes menos ruído. Com \(K = 1\) uma única estimativa do gradiente pode até ter o sinal errado e o treino ainda funciona: o SGD só precisa de gradientes certos em média e um lote de 128 imagens já tira a média de 128 sorteios independentes de \(\epsilon\). É por isso que todo VAE usa uma amostra por imagem.

Passo 3 — ajustando q com o truque

Agora treine \(q\). Comece no prior (\(m = 0\), \(\log v = 0\)) e suba o ELBO por subida estocástica do gradiente, uma amostra de \(\epsilon\) por passo, com os dois gradientes do truque da reparametrização. As curvas de nível são o ELBO exato; a linha branca é o que o SGD enxerga.

"""Lab 3 — fitting q by stochastic gradient ascent, one sample of z per step.

Same toy model (p(z) = N(0,1), p(x|z) = N(2z, 1), x = 3). q = N(m, e^s)
starts at the prior (m = 0, s = 0) and climbs the ELBO with the
reparameterization trick and K = 1, exactly as a VAE encoder would, but for a
single x. The contours are the exact ELBO; the path is what SGD sees.
"""

from io import StringIO

import matplotlib.pyplot as plt
import numpy as np

W, S, X = 2.0, 1.0, 3.0
VAR_X = W**2 + S**2
M_POST, V_POST = W * X / VAR_X, S**2 / VAR_X


def elbo(m, s):
    v = np.exp(s)
    rec = -0.5 * np.log(2 * np.pi * S**2) - ((X - W * m) ** 2 + W**2 * v) / (2 * S**2)
    return rec - 0.5 * (m**2 + v - 1 - s)


rng = np.random.default_rng(42)
m, s, lr, steps = 0.0, 0.0, 0.02, 1500
path = [(m, s)]
for _ in range(steps):
    eps = rng.standard_normal()
    sigma = np.exp(s / 2)
    z = m + sigma * eps
    dz = W * (X - W * z) / S**2                        # d log p(x|z) / dz
    gm = dz * 1.0 - m                                  # reconstruction + KL parts
    gs = dz * 0.5 * sigma * eps - 0.5 * (np.exp(s) - 1)
    m, s = m + lr * gm, s + lr * gs                    # ascent: maximize the ELBO
    path.append((m, s))
path = np.array(path)
tail = path[-500:].mean(0)

mm, ss = np.meshgrid(np.linspace(-0.5, 2.5, 300), np.linspace(-6, 1, 300))
fig, ax = plt.subplots(figsize=(7.5, 4.6))
cs = ax.contourf(mm, ss, elbo(mm, ss), levels=np.linspace(-12, -2.6, 25), cmap="viridis", extend="min")
cs.set_rasterized(True)
fig.colorbar(cs, ax=ax, label="ELBO", ticks=[-12, -10, -8, -6, -4, -2.62])
ax.plot(path[:, 0], path[:, 1], c="w", lw=0.8, alpha=0.9)
ax.plot(*path[0], "o", c="C0", ms=9, mec="w", label="start: $p(z) = \\mathcal{N}(0, 1)$")
ax.plot(*tail, "o", mfc="none", mec="C3", mew=2.2, ms=17,
        label=f"SGD, mean of last 500 steps: $m = {tail[0]:.2f}$, $v = {np.exp(tail[1]):.2f}$")
ax.plot(M_POST, np.log(V_POST), "*", c="k", ms=14, mec="w", zorder=5,
        label=f"$p(z\\mid x) = \\mathcal{{N}}({M_POST:.1f},\\ {V_POST:.1f})$")
ax.plot(1.5, np.log(0.01), "s", c="C1", ms=8, mec="w", label="$q_B = \\mathcal{N}(1.5,\\ 0.1^2)$")
ax.set_xlabel("$m$ (mean of $q$)")
ax.set_ylabel("$s = \\log v$ (log-variance of $q$)")
ax.set_title(f"{steps} steps, $K = 1$, $\\eta = {lr}$")
ax.legend(loc="lower left", fontsize=8.5, framealpha=0.85)
fig.tight_layout()
buf = StringIO()
fig.savefig(buf, format="svg", transparent=True, dpi=110)
print(buf.getvalue())

O caminho sobe direto pela encosta e depois vagueia numa nuvem pequena em volta do ótimo, porque cada passo vê um só \(\epsilon\). Na média dos últimos 500 passos ele fica em \(m = 1.20\), \(v = 0.20\): a posterior verdadeira, encontrada sem nunca calcular \(p(x)\) nem \(p(z \mid x)\). Só a KL o teria mantido no prior; só a reconstrução o teria levado até o quadrado \(q_B\).

Um encoder de verdade faz a mesma coisa para todas as imagens ao mesmo tempo, com uma diferença importante: em vez de números livres \((m, s)\) por imagem, ele aprende uma função \(x \mapsto (\mu, \log\sigma^2)\) compartilhada por todas. Essa função não consegue acertar cada posterior exatamente e a sobra se chama folga de amortização (amortization gap).

Passo 4 — o β move o ótimo

Troque o ELBO pelo objetivo ponderado por \(\beta\), \(\mathbb{E}_q[\log p(x \mid z)] - \beta\,\text{KL}(q \,\|\, p(z))\). Neste modelo o melhor chute tem fórmula exata, obtida igualando as duas derivadas a zero:

As cinco primeiras colunas são para \(x = 3\). A última tira a média da KL sobre um conjunto de 10 000 observações \(x \sim p(x)\): é quantos nats sobre \(x\) o código guarda.

"""Lab 4 — what beta does to the best q, in closed form.

Same toy model (p(z) = N(0,1), p(x|z) = N(w z, s^2), w = 2, s = 1). Maximizing
E_q[log p(x|z)] - beta * KL(q || p(z)) over q = N(m, v) has an exact answer:

    m* = w x / (w^2 + beta s^2)        v* = beta s^2 / (w^2 + beta s^2)

The first block is the single observation x = 3. The second averages over a
whole dataset x ~ p(x): the mean KL(q || p(z)) is how many nats about x the
code keeps.

Printed as a markdown table, in identifiers only, so one artifact serves both
the English and the Portuguese page.
"""

import numpy as np

W, S, X = 2.0, 1.0, 3.0
BETAS = (0.01, 0.1, 1.0, 4.0, 16.0, 100.0)


def best_q(x, beta):
    m = W * x / (W**2 + beta * S**2)
    v = beta * S**2 / (W**2 + beta * S**2)
    return m, v


def terms(x, m, v):
    rec = -0.5 * np.log(2 * np.pi * S**2) - ((x - W * m) ** 2 + W**2 * v) / (2 * S**2)
    kl = 0.5 * (m**2 + v - 1 - np.log(v))
    return rec, kl


rng = np.random.default_rng(42)
data = W * rng.standard_normal(10_000) + S * rng.standard_normal(10_000)   # x ~ p(x)

print("| `beta` | `m*` | `v*` | `E_q[log p(x∣z)]` | `KL(q ‖ p(z))` | `mean KL over the dataset` |")
print("|---:|---:|---:|---:|---:|---:|")
for beta in BETAS:
    m, v = best_q(X, beta)
    rec, kl = terms(X, m, v)
    dm, dv = best_q(data, beta)
    mean_kl = terms(data, dm, dv)[1].mean()
    bold = "**" if beta == 1.0 else ""
    print(f"| {bold}{beta:g}{bold} | {m:.3f} | {v:.4f} | {rec:.3f} | {kl:.3f} | {mean_kl:.3f} |")

Os três regimes de O β é a história inteira, em números:

• \(\beta \to 0\): \(v^* \to 0\) e \(m^* \to x/2 = 1.5\) — o ponto \(q_B\). A reconstrução é a melhor possível (\(-0.92\)) e o código guarda mais de 3 nats por observação. Um autoencoder.
• \(\beta = 1\): \(\mathcal{N}(1.2,\ 0.2)\), a posterior verdadeira. O código guarda 0.805 nats em média — exatamente \(\tfrac12\log 5\), a informação mútua entre \(x\) e \(z\) neste modelo.
• \(\beta \to \infty\): \(m^* \to 0\) e \(v^* \to 1\) — o prior. Em \(\beta = 100\) o código guarda 0.001 nat: ele não diz nada sobre \(x\). Colapso do posterior.

Aqui o decoder é fixo, então o colapso é só o encoder desistindo. Num VAE de verdade o decoder também aprende e isso piora as coisas: quando os códigos não carregam nada, o decoder aprende a ignorá-los — e aí, como o Experimente do passo 1 mostrou, o prior passa a ser mesmo a melhor resposta. Esse laço de realimentação é por que é difícil sair do colapso depois que ele começa e por que o annealing da KL começa com \(\beta\) em zero.

Adicional

Relação entre Log Variância e Desvio Padrão

Apêndice: a divergência KL

A divergência de Kullback–Leibler aparece três vezes neste capítulo — como a folga do ELBO, como o termo que puxa cada código na direção do prior e no argumento da KL para frente que explica o borrão. Vale conhecê-la por si só.

O que ela mede

Leia como um procedimento: sorteie \(x\) de \(q\), pergunte o quanto \(q\) o acha mais provável do que \(p\) acha (em escala log) e tire a média. Em termos de informação, é o número de nats extras por amostra que se paga quando dados que vêm de fato de \(q\) são codificados com um código feito para \(p\).

Três propriedades, todas visíveis no painel abaixo:

1. Ela nunca é negativa e só é zero quando \(q = p\) (desigualdade de Gibbs).
2. Ela não é simétrica: em geral \(\text{KL}(q \,\|\, p) \neq \text{KL}(p \,\|\, q)\), então não é uma distância.
3. O integrando pode ser negativo onde \(p > q\), mas o total nunca é.

Entre duas gaussianas

Para duas gaussianas a integral tem forma fechada:

Cada pedaço tem um significado: a razão entre as duas larguras, a dispersão de \(q\) medida em unidades de \(p\) e a distância ao quadrado entre as médias, também em unidades de \(p\). Faça \(p = \mathcal{N}(0, 1)\) e ela vira \(\tfrac12\big(\mu^2 + \sigma^2 - 1 - \log \sigma^2\big)\) — o termo KL da perda do VAE, deduzido no passo 5 a partir da esperança e aqui lido direto da fórmula geral.

Mova você mesmo

O que experimentar:

• Posterior do VAE vs prior coloca \(q = \mathcal{N}(1.2,\ 0.2)\), a posterior exata do exemplo resolvido, contra o prior. \(\text{KL}(q \,\|\, p) = 1.1247\): a coluna KL da tabela daquele exemplo.
• q estreita dentro de p dá \(\text{KL}(q \,\|\, p) = 0.318\) mas \(\text{KL}(p \,\|\, q) = 0.807\). Uma \(q\) estreita parada onde \(p\) tem massa é barata numa direção e cara na outra. Essa assimetria é a história de busca de modo contra cobertura de modos do capítulo 16.
• Bem afastadas coloca as médias em \(\pm 2\) com larguras iguais: as duas direções dão 8, porque com dispersões iguais a KL é simétrica, \((\mu_q - \mu_p)^2 / 2\sigma^2\). Arraste uma das larguras e veja os dois números se separarem.
• Aperte \(\sigma_q\) até 0.25 em cima de \(p = \mathcal{N}(0, 1)\): \(\text{KL}(q \,\|\, p)\) cresce como \(\log(1/\sigma_q)\). É o preço que um VAE paga por um encoder que vira autoencoder, um ponto nítido por entrada.