Diffusion math: score matching e SDE/ODE

Você pode rodar Stable Diffusion ou Flux sem entender uma única equação. O resultado vai sair, e até bonito. Mas no momento em que algo dá errado — imagem borrada, prompt ignorado, CFG instável, sampler ruim — você fica refém de chutar valores no cfg_scale e trocar de sampler até funcionar. Saber a matemática transforma esse processo num diagnóstico: por que 50 passos de Euler com CFG 12 saturam? Por que DPM-Solver++ 2M converge melhor que DDPM em 20 passos? Por que a probability flow ODE permite editar latentes de forma determinística?

Este módulo cobre o núcleo: forward gaussian process, reverse process, score function, equivalência ruído↔score, DDPM (Ho et al. 2020), DDIM (Song et al. 2021), e a formulação SDE/ODE contínua (Song, Sohl-Dickstein et al. 2021). É o background mínimo para ler papers de SD3, Flux, DiT e entender debates atuais sobre Flow Matching e Rectified Flow.

A grande sacada do DDPM (Ho, Jain, Abbeel — NeurIPS 2020) foi mostrar que o reverso, apesar de envolver uma integral intratável sobre todos os x₀ possíveis, pode ser aproximado por uma rede neural treinada com uma loss simples de MSE sobre o ruído adicionado. Sem adversarial training, sem reparametrization complicada, sem balanceamento delicado de GANs.

Defina um schedule de variâncias β₁, ..., β_T ∈ (0, 1) (linear, cosine, sigmoide — qualquer um monotônico crescente). A cada passo t você adiciona ruído gaussiano:

Aplicando recursivamente e usando linearidade de gaussianas, dá pra pular direto de x₀ até x_t em uma única expressão fechada — essencial para treinar eficientemente:

O verdadeiro reverso q(x_{t-1} | x_t) é intratável (depende de p(x_0)). Mas condicionado em x_0 ele é gaussiano e analítico:

q(x_{t-1} | x_t, x_0) = N(x_{t-1}; μ̃_t(x_t, x_0), β̃_t · I)

μ̃_t(x_t, x_0) = (√ᾱ_{t-1} · β_t / (1 − ᾱ_t)) · x_0
              + (√α_t · (1 − ᾱ_{t-1}) / (1 − ᾱ_t)) · x_t

β̃_t = ((1 − ᾱ_{t-1}) / (1 − ᾱ_t)) · β_t

O modelo p_θ(x_{t-1} | x_t) aproxima esse alvo. Em vez de aprender μ̃ diretamente, Ho et al. mostraram que prever o ruído ε é mais estável. Substituindo x_0 = (x_t − √(1 − ᾱ_t) · ε) / √ᾱ_t na expressão de μ̃, chega-se à parametrização ε-prediction:

A ELBO (Evidence Lower Bound) original do DDPM se decompõe em termos KL por timestep. Ho et al. mostraram que a versão simplificada(ignorando pesos do KL e usando apenas MSE no ruído) treina melhor na prática:

Em pseudocódigo Pythonesco — quase tão simples quanto treinar um classificador:

# DDPM training step (Ho et al. 2020, Algorithm 1)
def train_step(x0, model, alphas_bar, T):
    B = x0.shape[0]
    t = torch.randint(0, T, (B,), device=x0.device)        # timestep uniforme
    eps = torch.randn_like(x0)                              # ruído gaussiano
    a_bar = alphas_bar[t].view(B, 1, 1, 1)                 # ᾱ_t
    x_t = a_bar.sqrt() * x0 + (1 - a_bar).sqrt() * eps      # forward "fechado"
    eps_pred = model(x_t, t)                                # ε_θ(x_t, t)
    loss = F.mse_loss(eps_pred, eps)                        # L_simple
    return loss

Custo por imagem: T forwards da U-Net (~1000 no paper original). Esse é o gargalo prático que motivou DDIM.

Hyvärinen 2005 propôs estimar densidades aprendendo a score function s(x) = ∇_x log p(x). A intuição física: a score é um campo vetorial que aponta para regiões de maior densidade — se você seguir o gradiente em pequenos passos com ruído (Langevin), eventualmente cai numa amostra típica da distribuição.

Song & Ermon (NeurIPS 2019, NCSN — Noise Conditional Score Networks) notaram um problema: regiões de baixa densidade têm score mal estimada. A solução: perturbe os dados com múltiplos níveis de ruído e treine uma score-network condicional s_θ(x, σ) que estima a score em cada nível. Comece amostrando com ruído alto (cobre todo o espaço) e vá diminuindo σ progressivamente — annealed Langevin dynamics.

Song, Meng, Ermon (ICLR 2021) reformularam o reverso como um processo não-Markoviano: o estado x_{t-1} depende explicitamente de x_t e de x_0 predito, sem precisar passar por todos os intermediários. Isso quebra a obrigatoriedade do schedule original e permite saltar passos.

x_{t-1} = √ᾱ_{t-1} · x̂_0(x_t)
        + √(1 − ᾱ_{t-1} − σ_t²) · ε_θ(x_t, t)
        + σ_t · z

x̂_0(x_t) = (x_t − √(1 − ᾱ_t) · ε_θ(x_t, t)) / √ᾱ_t

σ_t = η · √( (1 − ᾱ_{t-1}) / (1 − ᾱ_t) ) · √(1 − α_t / α_{t-1})

O passo conceitual mais importante depois do DDPM. Song, Sohl-Dickstein, Kingma, Kumar, Ermon & Poole (ICLR 2021) mostraram que tudo isso é a discretização de uma SDE estocástica contínua no tempo:

Duas escolhas canônicas:

Pelo teorema de Anderson (1982), toda SDE forward admite uma SDE reversa com score:

E o pulo do gato: existe uma probability flow ODE determinística com mesma distribuição marginal: