Stable Diffusion 3.5 e Flux: MMDiT e DiT por dentro

De 2020 a 2023 a U-Net foi a arquitetura padrão. Funcionou em SD 1.5 (859M), SD 2.x e até SDXL (2.6B). Mas a partir de ~3B de parâmetros, U-Net começou a apresentar diminishing returns: mais compute não melhorava FID na mesma taxa.

Em janeiro de 2023, Peebles & Xie (ICCV — "Scalable Diffusion Models with Transformers") publicaram o DiT (Diffusion Transformer). A tese era simples: substituir completamente U-Net por um transformer puro patchificando o latente, e a escala se comportaria como NLP.

U-Net injeta timestep via FiLM em GroupNorm. DiT usa adaLN-Zero (Adaptive Layer Norm com pesos inicializados em zero):

class DiTBlock(nn.Module):
    def __init__(self, dim, heads, mlp_ratio=4):
        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False)
        self.attn  = MultiHeadAttention(dim, heads)
        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False)
        self.mlp   = FeedForward(dim, mlp_ratio * dim)

        # adaLN-Zero: 6 modulações (γ, β, α) × 2 (attn, mlp)
        self.modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(dim, 6 * dim, bias=True),
        )
        # CRÍTICO: pesos zerados na inicialização → bloco vira identidade no início
        nn.init.zeros_(self.modulation[-1].weight)
        nn.init.zeros_(self.modulation[-1].bias)

    def forward(self, x, c):     # c = t_emb + class/text_emb
        shift1, scale1, gate1, shift2, scale2, gate2 = self.modulation(c).chunk(6, dim=-1)
        # Attention com modulação
        x = x + gate1.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift1, scale1))
        # MLP com modulação
        x = x + gate2.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift2, scale2))
        return x

def modulate(x, shift, scale):
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

Esser, Kulal, Lorenz et al. 2024 (paper técnico do Stable Diffusion 3) propuseram uma evolução: dois streams paralelos de tokens — um para imagem, outro para texto — que se misturam via joint attention em cada bloco.

SD3 e Flux abandonaram o ε-prediction do DDPM em favor de Rectified Flow (Liu et al. 2022) / Flow Matching (Lipman et al. 2023). A intuição:

Em vez de pensar em "ruído sendo removido", pense em um fluxo que transporta uma amostra de ruído z₁ ~ N(0, I) até uma amostra de dado z₀ ~ p_data ao longo do tempo t ∈ [0, 1]. O caminho mais simples é a interpolação linear:

O vector field alvo ao longo dessa trajetória é simplesmente v(z_t, t) = z_1 − z_0 (derivada da interpolação). O modelo aprende:

Para amostrar, integra-se a ODE inversa partindo de z_1 ~ N(0, I):

def sample(model, n_steps=20):
    z = torch.randn(1, 4, 64, 64)               # z_1 ~ N(0, I)
    timesteps = torch.linspace(1, 0, n_steps + 1)
    for i in range(n_steps):
        t = timesteps[i]; t_next = timesteps[i + 1]
        v = model(z, t)                          # predição de velocidade
        z = z + (t_next - t) * v                 # Euler step
    return z                                     # z ≈ z_0 (latente da imagem)

Detalhe importante: o latente 64×64×4 vira sequência de tokens via patchify. Patches 2×2 produzem 32×32 = 1024 tokens de dim 16 (=2·2·4). DiT-XL/2 usa patches 2×2 (n_patches=256 num latente 32×32). SD3 usa esquema similar mas com patches menores em alta resolução.

Position embeddings: DiT original usa sinusoidal 2D (mesma do ViT). MMDiT e Flux usam RoPE 2D (Rotary Position Embedding extendido para grade), seguindo Su et al. 2021. RoPE é mais bem-comportado para extrapolação para resoluções maiores que o treino.