Phase 04 - Lesson 10

Geração de Imagens — Modelos de Difusão

Um modelo de difusão aprende a remover ruído. Treine-o para remover uma pequena quantidade de ruído de uma imagem ruidosa, repita isso para trás mil vezes e você tem um gerador de imagens.

Tipo: Build Linguagens: Python Pré-requisitos: Fase 4 Lição 07 (U-Net), Fase 1 Lição 06 (Probabilidade), Fase 3 Lição 06 (Otimizadores) Tempo: ~75 minutos

Objetivos de Aprendizagem

• Derivar o processo de adição de ruído direto x_0 -> x_1 -> ... -> x_T e explicar por que a forma fechada q(x_t | x_0) vale para qualquer t
• Implementar um objetivo de treino no estilo DDPM que faz regressão do ruído adicionado em cada passo, e um amostrador que caminha de volta do ruído puro até uma imagem
• Construir uma U-Net condicionada ao tempo (pequena o suficiente para treinar em CPU) que prevê o ruído para qualquer timestep
• Explicar a diferença entre amostragem DDPM e DDIM, e quando cada uma é apropriada (a Lição 23 cobre flow matching e rectified flow em profundidade)

O Problema

GANs geram em um único disparo: ruído entra, imagem sai, uma passagem para frente. São rápidas e difíceis de treinar. Modelos de difusão geram de forma iterativa: começam do ruído puro, removem ruído em pequenos passos, a imagem emerge. São lentos e fáceis de treinar. Nos últimos cinco anos essa última propriedade dominou: qualquer equipe pequena consegue treinar um modelo de difusão e obter amostras razoáveis; treinar uma GAN é um ofício que se aprende ao longo de anos de execuções fracassadas.

Além da estabilidade de treino, a estrutura iterativa da difusão é o que destrava tudo o que a geração moderna de imagens faz: condicionamento por texto, inpainting, edição de imagens, super-resolução, estilo controlável. Cada passo do laço de amostragem é um lugar para injetar uma nova restrição. Esse gancho é o motivo pelo qual Stable Diffusion, Imagen, DALL-E 3, Midjourney e todo modelo de imagem controlável que você vai usar são todos baseados em difusão.

Esta lição constrói o DDPM mínimo: adição de ruído direta, remoção de ruído reversa, laço de treino. A próxima lição (Stable Diffusion) o conecta a um sistema de produção com um VAE, um codificador de texto e classifier-free guidance.

O Conceito

O processo direto

Pegue uma imagem x_0. Adicione uma pequena quantidade de ruído gaussiano para obter x_1. Adicione um pouco mais para obter x_2. Continue por T passos até que x_T seja quase indistinguível de ruído gaussiano puro.

q(x_t | x_{t-1}) = N(x_t; sqrt(1 - beta_t) * x_{t-1},  beta_t * I)

beta_t é uma pequena agenda de variância, tipicamente linear de 0.0001 a 0.02 ao longo de T=1000 passos. Cada passo encolhe levemente o sinal e injeta ruído novo.

O salto em forma fechada

Adicionar ruído um passo de cada vez é uma cadeia de Markov, mas a matemática se condensa: você pode amostrar x_t diretamente a partir de x_0 em um único passo.

Define alpha_t = 1 - beta_t
Define alpha_bar_t = prod_{s=1..t} alpha_s

Then:
  q(x_t | x_0) = N(x_t; sqrt(alpha_bar_t) * x_0,  (1 - alpha_bar_t) * I)

Equivalently:
  x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon
  where epsilon ~ N(0, I)

Essa única equação é toda a razão pela qual a difusão é prática. Durante o treino você escolhe um t aleatório, amostra x_t diretamente de x_0 e treina em um único passo — sem necessidade de simular toda a cadeia de Markov.

O processo reverso

O processo direto é fixo. O processo reverso p(x_{t-1} | x_t) é o que a rede neural aprende. Modelos de difusão não preveem x_{t-1} diretamente; eles preveem o ruído epsilon adicionado no passo t, e a matemática deriva x_{t-1} a partir dele.

flowchart LR
    X0["x_0<br/>(imagem limpa)"] --> Q1["q(x_t|x_0)<br/>adiciona ruído"]
    Q1 --> XT["x_t<br/>(ruidosa)"]
    XT --> MODEL["model(x_t, t)"]
    MODEL --> EPS["epsilon previsto"]
    EPS --> LOSS["MSE contra<br/>epsilon verdadeiro"]

    XT -.->|amostragem| STEP["p(x_{t-1}|x_t)"]
    STEP -.-> XT1["x_{t-1}"]
    XT1 -.->|repete 1000x| X0S["x_0 (amostrado)"]

    style X0 fill:#dcfce7,stroke:#16a34a
    style MODEL fill:#fef3c7,stroke:#d97706
    style LOSS fill:#fecaca,stroke:#dc2626
    style X0S fill:#dbeafe,stroke:#2563eb

A função de perda do treino

Para cada passo de treino:

1. Amostre uma imagem real x_0.
2. Amostre um timestep t uniformemente de [1, T].
3. Amostre ruído epsilon ~ N(0, I).
4. Calcule x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon.
5. Preveja epsilon_theta(x_t, t) com a rede.
6. Minimize || epsilon - epsilon_theta(x_t, t) ||^2.

É isso. A rede neural aprende a prever o ruído em qualquer timestep. A perda é MSE. Não há jogo adversarial, nem colapso, nem oscilação.

O amostrador (DDPM)

Para gerar: comece de x_T ~ N(0, I) e caminhe para trás um passo de cada vez.

for t = T, T-1, ..., 1:
    eps = model(x_t, t)
    x_{t-1} = (1 / sqrt(alpha_t)) * (x_t - (beta_t / sqrt(1 - alpha_bar_t)) * eps) + sqrt(beta_t) * z
    where z ~ N(0, I) if t > 1, else 0
return x_0

O ponto chave é que, embora a condicional reversa não seja conhecida em forma fechada no caso geral, para este processo direto gaussiano específico ela é. Os coeficientes de aparência feia são o que a regra de Bayes lhe dá.

Por que 1000 passos

A agenda de ruído direto é escolhida de modo que cada passo adicione apenas ruído suficiente para que o passo reverso seja quase gaussiano. Poucos passos demais e o passo reverso fica longe de gaussiano, a rede não consegue modelá-lo bem. Passos demais e a amostragem fica cara com ganho decrescente. T=1000 com uma agenda linear é o padrão do DDPM.

DDIM: amostragem 20x mais rápida

O treino é o mesmo. A amostragem muda. O DDIM (Song et al., 2020) define um processo reverso determinístico que pula timesteps sem retreinar. Amostrar em 50 passos com DDIM dá qualidade próxima à do DDPM de 1000 passos. Todo sistema de produção usa DDIM ou uma variante ainda mais rápida (DPM-Solver, Euler ancestral).

Condicionamento ao tempo

A rede epsilon_theta(x_t, t) precisa saber qual timestep ela está removendo o ruído. Modelos de difusão modernos injetam t via embeddings de tempo sinusoidais (mesma ideia da codificação posicional em transformers) que são adicionados aos mapas de características em cada nível da U-Net.

t_embedding = sinusoidal(t)
feature_map += MLP(t_embedding)

Sem condicionamento ao tempo a rede tem que adivinhar o nível de ruído a partir da própria imagem, o que funciona mas é muito menos eficiente em amostras.

Construa

Passo 1: Agenda de ruído

import torch

def linear_beta_schedule(T=1000, beta_start=1e-4, beta_end=2e-2):
    return torch.linspace(beta_start, beta_end, T)


def precompute_schedule(betas):
    alphas = 1.0 - betas
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    return {
        "betas": betas,
        "alphas": alphas,
        "alphas_cumprod": alphas_cumprod,
        "sqrt_alphas_cumprod": torch.sqrt(alphas_cumprod),
        "sqrt_one_minus_alphas_cumprod": torch.sqrt(1.0 - alphas_cumprod),
        "sqrt_recip_alphas": torch.sqrt(1.0 / alphas),
    }

schedule = precompute_schedule(linear_beta_schedule(T=1000))

Pré-compute uma vez, busque por índice durante o treino e a amostragem.

Passo 2: Difusão direta (q_sample)

def q_sample(x0, t, noise, schedule):
    sqrt_a = schedule["sqrt_alphas_cumprod"][t].view(-1, 1, 1, 1)
    sqrt_one_minus_a = schedule["sqrt_one_minus_alphas_cumprod"][t].view(-1, 1, 1, 1)
    return sqrt_a * x0 + sqrt_one_minus_a * noise

Forma fechada de uma linha. t é um lote de timesteps, um por imagem no lote.

Passo 3: Uma pequena U-Net condicionada ao tempo

import torch.nn as nn
import torch.nn.functional as F
import math

def timestep_embedding(t, dim=64):
    half = dim // 2
    freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half)
    args = t[:, None].float() * freqs[None]
    emb = torch.cat([args.sin(), args.cos()], dim=-1)
    return emb


class TinyUNet(nn.Module):
    def __init__(self, img_channels=3, base=32, t_dim=64):
        super().__init__()
        self.t_mlp = nn.Sequential(
            nn.Linear(t_dim, base * 4),
            nn.SiLU(),
            nn.Linear(base * 4, base * 4),
        )
        self.t_dim = t_dim
        self.enc1 = nn.Conv2d(img_channels, base, 3, padding=1)
        self.enc2 = nn.Conv2d(base, base * 2, 4, stride=2, padding=1)
        self.mid = nn.Conv2d(base * 2, base * 2, 3, padding=1)
        self.dec1 = nn.ConvTranspose2d(base * 2, base, 4, stride=2, padding=1)
        self.dec2 = nn.Conv2d(base * 2, img_channels, 3, padding=1)
        self.time_proj = nn.Linear(base * 4, base * 2)

    def forward(self, x, t):
        t_emb = timestep_embedding(t, self.t_dim)
        t_emb = self.t_mlp(t_emb)
        t_proj = self.time_proj(t_emb)[:, :, None, None]

        h1 = F.silu(self.enc1(x))
        h2 = F.silu(self.enc2(h1)) + t_proj
        h3 = F.silu(self.mid(h2))
        d1 = F.silu(self.dec1(h3))
        d2 = torch.cat([d1, h1], dim=1)
        return self.dec2(d2)

U-Net de dois níveis com condicionamento ao tempo injetado no gargalo. Aumente a profundidade e a largura para imagens reais.

Passo 4: Laço de treino

def train_step(model, x0, schedule, optimizer, device, T=1000):
    model.train()
    x0 = x0.to(device)
    bs = x0.size(0)
    t = torch.randint(0, T, (bs,), device=device)
    noise = torch.randn_like(x0)
    x_t = q_sample(x0, t, noise, schedule)
    pred = model(x_t, t)
    loss = F.mse_loss(pred, noise)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

Esse é o laço de treino inteiro. Sem jogo de GAN, sem perda especializada, uma única chamada de MSE.

Passo 5: Amostrador (DDPM)

@torch.no_grad()
def sample(model, schedule, shape, T=1000, device="cpu"):
    model.eval()
    x = torch.randn(shape, device=device)
    betas = schedule["betas"].to(device)
    sqrt_one_minus_a = schedule["sqrt_one_minus_alphas_cumprod"].to(device)
    sqrt_recip_alphas = schedule["sqrt_recip_alphas"].to(device)

    for t in reversed(range(T)):
        t_batch = torch.full((shape[0],), t, dtype=torch.long, device=device)
        eps = model(x, t_batch)
        coef = betas[t] / sqrt_one_minus_a[t]
        mean = sqrt_recip_alphas[t] * (x - coef * eps)
        if t > 0:
            x = mean + torch.sqrt(betas[t]) * torch.randn_like(x)
        else:
            x = mean
    return x

1000 passagens para frente para produzir um lote de amostras. Em código real você trocaria isso por um amostrador DDIM de 50 passos.

Passo 6: Amostrador DDIM (determinístico, ~20x mais rápido)

@torch.no_grad()
def sample_ddim(model, schedule, shape, steps=50, T=1000, device="cpu", eta=0.0):
    model.eval()
    x = torch.randn(shape, device=device)
    alphas_cumprod = schedule["alphas_cumprod"].to(device)

    ts = torch.linspace(T - 1, 0, steps + 1).long()
    for i in range(steps):
        t = ts[i]
        t_prev = ts[i + 1]
        t_batch = torch.full((shape[0],), t, dtype=torch.long, device=device)
        eps = model(x, t_batch)
        a_t = alphas_cumprod[t]
        a_prev = alphas_cumprod[t_prev] if t_prev >= 0 else torch.tensor(1.0, device=device)
        x0_pred = (x - torch.sqrt(1 - a_t) * eps) / torch.sqrt(a_t)
        sigma = eta * torch.sqrt((1 - a_prev) / (1 - a_t) * (1 - a_t / a_prev))
        dir_xt = torch.sqrt(1 - a_prev - sigma ** 2) * eps
        noise = sigma * torch.randn_like(x) if eta > 0 else 0
        x = torch.sqrt(a_prev) * x0_pred + dir_xt + noise
    return x

eta=0 é totalmente determinístico (a mesma entrada de ruído sempre produz a mesma saída). eta=1 recupera o DDPM.

Use

Para trabalho de produção, use diffusers:

from diffusers import DDPMScheduler, UNet2DModel

unet = UNet2DModel(sample_size=32, in_channels=3, out_channels=3, layers_per_block=2)
scheduler = DDPMScheduler(num_train_timesteps=1000)

A biblioteca traz schedulers prontos (DDPM, DDIM, DPM-Solver, Euler, Heun), U-Nets configuráveis, pipelines para texto-para-imagem e imagem-para-imagem, e utilitários de fine-tuning com LoRA.

Para pesquisa, k-diffusion (Katherine Crowson) tem as implementações de referência mais fiéis e as melhores variantes de amostragem.

Entregue

Esta lição produz:

• outputs/prompt-diffusion-sampler-picker.md — um prompt que escolhe DDPM / DDIM / DPM-Solver / Euler com base na meta de qualidade, no orçamento de latência e no tipo de condicionamento.
• outputs/skill-noise-schedule-designer.md — uma skill que produz uma agenda beta linear, cosseno ou sigmoide dado T e o nível de corrupção alvo, além de gráficos de diagnóstico da razão sinal-ruído ao longo do tempo.

Exercícios

1. (Fácil) Visualize o processo direto: pegue uma imagem e plote x_t em t in [0, 100, 250, 500, 750, 1000]. Verifique que x_1000 parece ruído gaussiano puro.
2. (Médio) Treine a TinyUNet no dataset de círculos sintéticos por 20 épocas e amostre 16 círculos. Compare a amostragem DDPM (1000 passos) e DDIM (50 passos) — elas produzem imagens semelhantes a partir da mesma semente de ruído?
3. (Difícil) Implemente uma agenda de ruído por cosseno (Nichol & Dhariwal, 2021): alpha_bar_t = cos^2((t/T + s) / (1 + s) * pi / 2). Treine o mesmo modelo com agendas linear e cosseno e mostre que o cosseno dá melhores amostras com baixa contagem de passos.

Termos-Chave

Leitura Complementar

• Denoising Diffusion Probabilistic Models (Ho et al., 2020) — o artigo que tornou a difusão prática e superou as GANs em FID
• Improved DDPM (Nichol & Dhariwal, 2021) — agenda por cosseno e parametrização v
• DDIM (Song, Meng, Ermon, 2020) — o amostrador determinístico que tornou possível a inferência em tempo real
• Elucidating the Design Space of Diffusion (Karras et al., 2022) — uma visão unificada de cada escolha de design de difusão; a melhor referência atual