Phase 08 - Lesson 02

Autoencoders & Autoencoders Variacionais (VAE)

Um autoencoder comum comprime e depois reconstrói. Ele memoriza. Ele não gera. Adicione um truque — force o código a parecer gaussiano — e você terá um amostrador. Esse único truque, a reparametrização de z = μ + σ·ε, é o motivo pelo qual todo modelo de imagem por difusão latente e flow-matching que você usa em 2026 tem um VAE na entrada.

Tipo: Construir Linguagens: Python Pré-requisitos: Fase 3 · 02 (Backprop), Fase 3 · 07 (CNNs), Fase 8 · 01 (Taxonomia) Tempo: ~75 minutos

O Problema

Comprima um dígito MNIST de 784 pixels para um código de 16 números e depois o reconstrua. Um autoencoder comum terá um excelente desempenho no MSE de reconstrução, mas o espaço de código será uma bagunça irregular. Escolha um ponto aleatório no espaço de código, decodifique-o e você obterá ruído. Ele não tem amostrador. É apenas um modelo de compressão disfarçado.

O que você realmente quer é: (a) que o espaço de código seja uma distribuição limpa e suave da qual você possa amostrar — como uma Gaussiana isotrópica N(0, I), (b) que a decodificação de qualquer amostra produza um dígito plausível e (c) que o codificador e o decodificador ainda comprimam bem. Três objetivos, uma arquitetura, uma perda.

O VAE de Kingma (2013) resolve isso treinando o codificador para gerar uma distribuição q(z|x) = N(μ(x), σ(x)²), puxando essa distribuição em direção à priori N(0, I) por meio de uma penalidade KL e, em seguida, amostrando z de q(z|x) antes de decodificar. No momento da inferência, descarte o codificador, amostre z ~ N(0, I) e decodifique. A penalidade KL é o que força o espaço de código a ser estruturado.

Em 2026, os VAEs raramente são implantados de forma autônoma — eles foram superados pela difusão em termos de qualidade bruta de imagem —, mas são o codificador de escolha para todo modelo de difusão latente (SD 1/2/XL/3, Flux, AudioCraft). Aprenda o VAE e você aprenderá a primeira camada invisível de todo pipeline de imagem que utiliza.

O Conceito

Autoencoder. z = encoder(x), x̂ = decoder(z), loss = ||x - x̂||². Espaço de código não estruturado.

Codificador VAE. Gera dois vetores: μ(x) e log σ²(x). Eles definem q(z|x) = N(μ, diag(σ²)).

Truque de reparametrização. A amostragem de q(z|x) não é diferenciável. Reescreva a amostra como z = μ + σ·ε, onde ε ~ N(0, I). Agora z é uma função determinante de (μ, σ) mais um ruído não parametrizado — os gradientes fluem através de μ e σ.

Perda. Evidence Lower BOund (ELBO), dois termos:

loss = reconstruction + β · KL[q(z|x) || N(0, I)]
     = ||x - x̂||²  + β · Σ_i ( σ_i² + μ_i² - log σ_i² - 1 ) / 2

A reconstrução empurra x̂ em direção a x. O KL empurra q(z|x) em direção à priori. Eles se compensam (trade-off). β pequeno (<1) = amostras mais nítidas, espaço de código menos gaussiano. β grande (>1) = espaço de código mais limpo, amostras mais borradas. O β-VAE (Higgins 2017) tornou esse ajuste famoso e iniciou a pesquisa de desentrelaçamento (disentanglement).

Amostragem. Na inferência: extraia z ~ N(0, I), passe adiante pelo decodificador. Apenas um passo à frente (forward pass) — sem amostragem iterativa como na difusão.

Construa

code/main.py implementa um mini VAE sem numpy ou torch. A entrada é um dado sintético de 8 dimensões extraído de uma mistura gaussiana de 2 componentes em 8-D. O codificador e o decodificador são MLPs de uma única camada oculta. Implementamos a ativação tanh, o passo à frente (forward pass), a perda e um passo de retropropagação (backward pass) escrito à mão. Não é para produção — é pedagógico.

Passo 1: passo à frente do codificador

def encode(x, enc):
    h = tanh(add(matmul(enc["W1"], x), enc["b1"]))
    mu = add(matmul(enc["W_mu"], h), enc["b_mu"])
    log_sigma2 = add(matmul(enc["W_sig"], h), enc["b_sig"])
    return mu, log_sigma2

log σ² em vez de σ para que a saída da rede seja irrestrita (softplus de σ é uma armadilha — os gradientes morrem em σ ≈ 0).

Passo 2: reparametrizar e decodificar

def reparameterize(mu, log_sigma2, rng):
    eps = [rng.gauss(0, 1) for _ in mu]
    sigma = [math.exp(0.5 * lv) for lv in log_sigma2]
    return [m + s * e for m, s, e in zip(mu, sigma, eps)]

def decode(z, dec):
    h = tanh(add(matmul(dec["W1"], z), dec["b1"]))
    return add(matmul(dec["W_out"], h), dec["b_out"])

Passo 3: o ELBO

def elbo(x, x_hat, mu, log_sigma2, beta=1.0):
    recon = sum((a - b) ** 2 for a, b in zip(x, x_hat))
    kl = 0.5 * sum(math.exp(lv) + m * m - lv - 1 for m, lv in zip(mu, log_sigma2))
    return recon + beta * kl, recon, kl

KL de forma fechada exata porque ambas as distribuições são gaussianas. Não integre numericamente. As pessoas ainda enviam código com estimativas KL de monte-carlo em 2026 — é 3 vezes mais lento sem motivo.

Passo 4: gerar

def sample(dec, z_dim, rng):
    z = [rng.gauss(0, 1) for _ in range(z_dim)]
    return decode(z, dec)

Esse é o modelo gerativo. Cinco linhas.

Armadilhas

• Colapso do posterior (Posterior collapse). O termo KL direciona q(z|x) → N(0, I) tão agressivamente que z não carrega informações sobre x. Correção: recozimento de β (beta-annealing, comece com β=0 e aumente até 1), bits livres (free bits) ou ignore o KL em dimensões inativas.
• Amostras borradas. A verossimilhança do decodificador gaussiano implica em reconstrução por MSE, que é Bayes-otimizado para L2 (a média) — a média de um conjunto de dígitos plausíveis é um dígito difuso. Correção: decodificador discreto (VQ-VAE, NVAE) ou use o VAE apenas como um codificador e aplique difusão nas variáveis latentes (isso é o que o Stable Diffusion faz).
• β muito grande, muito cedo. Veja colapso do posterior. Comece em β≈0.01 e suba gradativamente.
• Dimensão latente muito pequena. 16-D funciona para o MNIST, 256-D para o ImageNet 256², 2048-D para o ImageNet 1024². O VAE do Stable Diffusion comprime 512×512×3 → 64×64×4 (fator de subamostragem de 32x em área espacial, 32x em canais).

Casos de Uso

A pilha VAE em 2026:

Um modelo de difusão latente é um VAE com um modelo de difusão situado entre o codificador e o decodificador. O VAE faz a compressão grosseira, o modelo de difusão faz o trabalho pesado. O mesmo padrão se aplica para vídeo (VAE + DiT de distribuição de vídeo) e áudio (Encodec + transformer MusicGen).

Envie

Salve outputs/skill-vae-trainer.md.

A habilidade (skill) recebe: perfil do dataset + meta de dimensão latente + uso posterior (reconstrução, amostragem ou entrada de difusão latente) e gera: escolha da arquitetura (comum/β/VQ/RVQ), cronograma de β, dimensão latente, verossimilhança do decodificador (gaussiano vs categórico) e plano de avaliação (MSE de reconstrução, KL por dimensão, distância de Fréchet entre q(z|x) e N(0, I)).

Exercícios

1. Fácil. Altere β em code/main.py para 0.01, 0.1, 1.0, 5.0. Registre o MSE de reconstrução final e o KL. Qual β é o melhor no sentido de Pareto (Pareto-best) para seus dados sintéticos?
2. Médio. Substitua a verossimilhança do decodificador gaussiano por uma verossimilhança de Bernoulli (perda de entropia cruzada). Compare a qualidade das amostras em uma versão binarizada dos mesmos dados sintéticos.
3. Difícil. Estenda o code/main.py em um mini VQ-VAE: substitua o z contínuo por uma busca pelo vizinho mais próximo em um livro de códigos (codebook) de K=32 entradas. Compare o MSE de reconstrução e relate quantas entradas do livro de códigos são realmente utilizadas (o colapso do livro de códigos é real).

Termos-Chave

Nota de produção: o VAE é o caminho mais crítico em um servidor de difusão

Em um pipeline do Stable Diffusion / Flux / SD3, o VAE é chamado duas vezes por requisição — uma para codificar (se estiver fazendo img2img / inpainting) e outra para decodificar. Em 1024², a passagem do decodificador é frequentemente o maior pico de memória de ativação em todo o pipeline porque ele faz o upsampling de latentes de 128×128×16 de volta para 1024×1024×3. Duas consequências práticas:

• Fatie ou divida em blocos (slice/tile) a decodificação. A biblioteca diffusers expõe pipe.vae.enable_slicing() e pipe.vae.enable_tiling(). O fatiamento em blocos (tiling) troca um pequeno artefato de costura por memória O(tile²) em vez de O(H·W). Essencial para 1024²+ em GPUs de consumo.
• Decodificador bf16, cálculos fp32 para o redimensionamiento final. O VAE do SD 1.x foi lançado em fp32 e produz NaNs silenciosamente quando convertido para fp16 em 1024²+. O SDXL distribui o madebyollin/sdxl-vae-fp16-fix — prefira sempre a variante fp16-fix ou use bf16.

Leituras Recomendadas

• Kingma & Welling (2013). Auto-Encoding Variational Bayes — o artigo do VAE.
• Higgins et al. (2017). β-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework — β-VAE desentrelaçado.
• van den Oord et al. (2017). Neural Discrete Representation Learning — VQ-VAE.
• Vahdat & Kautz (2021). NVAE: A Deep Hierarchical Variational Autoencoder — VAE de imagem no estado da arte.
• Rombach et al. (2022). High-Resolution Image Synthesis with Latent Diffusion Models — Stable Diffusion; VAE como codificador.
• Défossez et al. (2022). High Fidelity Neural Audio Compression — Encodec, o padrão de VAE de áudio.