22. Diffusion Transformers

Diffusion Transformers (DiT)

Em 2023, Peebles e Xie¹ pegaram o modelo de difusão latente do capítulo 20, tiraram a sua U-Net e puseram no lugar um Vision Transformer comum: cortar o latente ruidoso em patches, rodar blocos Transformer sobre os tokens, transformar os tokens de volta num latente. O resultado superou as melhores U-Nets da época no ImageNet condicionado por classe, mas o achado de verdade do artigo foi outro: a qualidade melhorava de forma previsível com o compute, do jeito que acontece com modelos de linguagem.

Três anos depois, essa troca venceu onde há dinheiro em jogo. Os principais modelos abertos de texto para imagem e de texto para vídeo de 2024–2026 — SD3, FLUX, Qwen-Image, Z-Image, HunyuanVideo, Wan, LTX-2 — descendem desse bloco, assim como os geradores dentro de políticas de robôs e de modelos de objetos 3-D. A principal exceção é a geração autorregressiva de imagens (capítulo 23); os sistemas fechados (Sora 2, Veo 3, Midjourney) não publicam a sua arquitetura.

Este capítulo desmonta o bloco: como uma imagem vira tokens, como o timestep entra (adaLN-Zero), como o prompt entra (cross-attention, atenção conjunta, um fluxo só), como as posições são codificadas quando a grade ganha um eixo de tempo e para onde vai o compute de fato — o que, no vídeo, é a história inteira.

Slides

Clique nos slides e use as setas ← → para navegar; F abre em tela cheia e P abre a visão do apresentador, com as notas. Os simuladores desta página rodam dentro dos slides. Abrir os slides numa aba nova

De U-Net para Transformer

O removedor de ruído do SD 1.x e do SDXL é uma U-Net: convoluções que reduzem a resolução à metade a cada nível e depois a trazem de volta, skip connections entre os níveis correspondentes e atenção — com o texto entrando por cross-attention — intercalada com as convoluções. Um DiT não tem nada dessa hierarquia. Ele trabalha numa única resolução do primeiro bloco ao último e todo bloco vê todos os tokens.

Todo o resto do pipeline fica: o VAE, o encoder de texto, o amostrador, o classifier-free guidance. Só muda a rede que prevê o ruído — ou, desde o SD3, a velocidade do flow matching.

Passo 1 — Patchify: um latente como sequência de tokens

Um DiT trabalha sobre o latente de um VAE, não sobre pixels. Um latente de formato \(C \times H \times W\) é cortado em patches \(p \times p\) sem sobreposição; cada patch, com \(p^2 C\) números, é achatado e projetado por uma camada linear para a largura \(d\) do modelo:

É exatamente o patch embedding do ViT — uma Conv2d com kernel e passo \(p\). A parte nova é o caminho de volta. Um gerador precisa devolver um latente do formato da entrada, então uma última camada linear leva cada token a \(p^2 C\) números e o unpatchify os recoloca na grade: um reshape e um permute, sem parâmetros.

O tamanho do patch é o principal botão de compute do DiT e o artigo o girou de propósito. Reduzir \(p\) à metade quadruplica \(N\) enquanto o número de parâmetros quase não se mexe (só \(E\) e a camada de saída dependem de \(p\)). Com 400K passos de treino, o DiT-XL com \(p = 8, 4, 2\) custa 7,4, 29,1 e 118,6 Gflops por passada (bilhões de operações, contando uma multiplicação-adição como uma) e chega a FID 106,4, 43,0 e 19,5¹ (a distância de Fréchet, nas features da Inception, entre 50 mil amostras e as imagens reais; menor é melhor): mesmos parâmetros, mais tokens, imagens muito melhores. É a observação sobre a qual se apoia todo o argumento de escala — o que prevê a qualidade é o compute por passada, venha ele da largura, da profundidade ou dos tokens. O laboratório reproduz a tendência em \(28 \times 28\).

Passo 2 — O bloco e como o timestep entra: adaLN-Zero

Um removedor de ruído precisa saber quão ruidosa está a entrada e, num modelo condicionado por classe, qual classe desenhar. O DiT junta as duas coisas num vetor,

em que \(\text{emb}(t)\) escreve \(t\) como senos e cossenos de muitas frequências (como nas posições das palavras) e \(\text{emb}(y)\) é a linha de uma tabela aprendida para a classe \(y\). Todo bloco lê \(c\) pelas suas camadas de normalização. Uma LayerNorm normalmente termina com uma escala e um deslocamento aprendidos, os mesmos para qualquer entrada. O adaLN (LayerNorm adaptativa) os remove e deixa uma camada linear prevê-los a partir de \(c\); o adaLN-Zero prevê também um portão \(\alpha\) para cada ramo residual. Uma camada linear por bloco produz seis vetores de tamanho \(d\):

Aqui \(x\) são os \(N\) tokens do bloco, de largura \(d\); \(\text{SiLU}(z) = z\,\sigma(z)\); LN é uma LayerNorm sem parâmetros próprios; \(\odot\) multiplica elemento a elemento; os mesmos seis vetores servem a todos os tokens da imagem. Percorra o bloco passo a passo: de onde vem \(c\), para onde vai cada um dos seis vetores e o que a inicialização faz com eles.

O "Zero" é a inicialização: \(W = 0\) e \(b = 0\), então no começo \(\alpha = 0\) e todo bloco é a identidade. A rede começa como uma pilha de operações nulas (a camada de saída também começa em zero, então ela prevê zero) e cada bloco aprende a acrescentar alguma coisa. É o mesmo truque de adição segura das conexões residuais do capítulo 9, das convoluções zeradas do ControlNet e do \(B = 0\) do LoRA.

O mesmo bloco como o laboratório o escreve, rodando em números pequenos o bastante para acompanhar: um latente \(1 \times 4 \times 4\) cortado em quatro patches \(2 \times 2\), \(d = 4\), uma linha do DiTBlock.forward por vez. Percorra o caso treinado; depois escolha na inicialização e veja os dois portões segurarem \(z\) parado:

O artigo do DiT testou quatro jeitos de entregar \(c\) ao bloco, com o mesmo orçamento de treino: no contexto (os embeddings de \(t\) e de \(y\) acrescentados à sequência como mais dois tokens), cross-attention (os tokens os leem numa camada de atenção extra), adaLN e adaLN-Zero.

O adaLN-Zero venceu com folga — FID 19,5 contra 35,2 dos tokens no contexto, com 400K passos — e com o menor compute (empatado com o adaLN), porque a modulação é calculada uma vez por imagem, não uma vez por token. O laboratório roda três dos quatro desenhos no Fashion-MNIST e encontra a mesma ordem: FD 26,62 para os tokens no contexto, 19,97 para o adaLN e 17,11 para o adaLN-Zero (o FD é a fórmula do FID calculada nas features de um classificador pequeno de Fashion-MNIST, no lugar das da Inception).

Passo 3 — Como o prompt entra: cross-attention, atenção conjunta, um fluxo só

O artigo do DiT tinha só um rótulo de classe. Um modelo de texto para imagem precisa de um prompt inteiro — uma sequência de centenas de embeddings de tokens vindos de um encoder de texto — e há três desenhos em uso.

Cross-attention. Os tokens de imagem fazem atenção entre si e depois leem o texto numa segunda camada de atenção, cujas chaves e valores vêm do prompt: \(\text{softmax}(Q_{\text{img}} K_{\text{txt}}^\top / \sqrt{d_h})\, V_{\text{txt}}\). É o que a U-Net do SD fazia e o que fazem o PixArt-α⁵, o Wan⁷ e o Movie Gen¹⁰. O texto é uma entrada só de leitura: ele sai de cada camada exatamente como entrou.

Atenção conjunta com dois fluxos (MMDiT, DiT multimodal). O SD3² concatena tokens de imagem e de texto numa única sequência e roda uma atenção sobre ela:

em que \([\,\cdot\,;\,\cdot\,]\) concatena ao longo da sequência e \(d_h\) é a largura de uma cabeça de atenção. Mas cada modalidade mantém os seus próprios pesos — as suas projeções Q, K, V, o seu MLP e a sua modulação. Os tokens de texto agora leem a imagem e são reescritos por ela, camada após camada. Na comparação do artigo do SD3, com o mesmo tamanho, o MMDiT venceu o DiT com cross-attention, um U-ViT (um ViT com skip connections longas entre os blocos rasos e os profundos) e um DiT comum que só acrescenta os tokens de texto à sequência.

Atenção conjunta com um fluxo só. Depois que as duas modalidades se alinharam por um tempo, os pesos separados podem sair: a sequência concatenada passa por um bloco Transformer comum, com um único conjunto de pesos. O FLUX.1 roda 19 blocos de dois fluxos e depois 38 de um fluxo³; o HunyuanVideo faz 20 e depois 40⁸; o FLUX.2 faz 8 e depois 48⁶; o Z-Image usa um fluxo só do primeiro bloco ao último²¹.

Percorra um bloco de dois fluxos e, no fim, o bloco de um fluxo só que o substitui:

O painel abaixo olha dentro da atenção: quais blocos da matriz de escores cada desenho calcula.

O painel deixa duas coisas visíveis. Primeiro, a atenção conjunta é quase de graça em compute: o texto é curto perto da imagem, então os blocos extras da matriz de escores acrescentam 12,5% aos escores da cross-attention com os 4096 + 512 tokens do FLUX e menos de 1% num vídeo em 720p. O que custa é o segundo conjunto de pesos dos blocos de dois fluxos — exatamente o que os blocos de um fluxo economizam. Segundo, a escolha não está decidida: o Wan mantém cross-attention com 14B de parâmetros, o Z-Image usa um fluxo só com 6B e os dois são competitivos.

Posições: de uma tabela de senos ao RoPE em duas e três dimensões

A atenção não enxerga ordem (capítulo 13), então os tokens precisam de posições. O DiT original soma aos tokens, depois do patchify, uma tabela fixa de senos e cossenos 2-D — metade dos canais codifica a linha, metade a coluna. Funciona na resolução de treino e precisa ser interpolada para qualquer outra.

Desde o FLUX.1, a maioria dos modelos novos usa no lugar dela rotary embeddings (RoPE²⁷) aplicadas por eixo (o SD3 ainda soma uma tabela 2-D). Os canais de cada cabeça de atenção são divididos em grupos e cada grupo gira \(q\) e \(k\) por um ângulo proporcional a uma coordenada; o escore \(q^\top k\) passa a depender só da diferença de posições ao longo de cada eixo.

As posições relativas são o que torna um DiT indiferente à resolução e à proporção: uma imagem de \(1024 \times 768\) é só uma grade diferente de coordenadas e um vídeo é a mesma grade com um terceiro eixo. Os tokens de texto também ganham posições — o Qwen-Image os coloca na diagonal da grade da imagem, para que os dois dividam um único sistema de coordenadas²⁰.

O que o DiT de fato comprou: uma lei de escala

A contribuição do artigo do DiT não é "um Transformer também funciona aqui". É que trocar o backbone deu à geração de imagens a propriedade que a modelagem de linguagem já tinha: uma curva com a qual se pode planejar¹. Em quatro tamanhos de modelo (S, B, L, XL) e três tamanhos de patch, o FID caiu de forma constante à medida que subiam os Gflops de uma passada, sem sinal de saturação na faixa testada. Os parâmetros sozinhos o previam mal: o DiT-XL/8, com 676M de parâmetros, terminou com FID 106,4, bem atrás do DiT-B/2 com 130M e 43,5; já o DiT-S/2 (33M) e o DiT-B/4 (130M), ambos com cerca de 6 Gflops, chegaram a 68,40 e 68,38. O SD3 repetiu o estudo na escala de texto para imagem (até 8B de parâmetros) e viu a loss de validação cair de forma suave com o tamanho e com os passos de treino, acompanhando o GenEval (um benchmark automático de aderência ao prompt: contagens, cores, posições) e a preferência humana².

Os doze modelos DiT do artigo, nos dois gráficos (troque o eixo horizontal):

Essa curva é o que justificou gastar na geração seguinte:

Duas coisas na tabela. O desenho do SD3 amarra a largura à profundidade (\(d = 64 \times\) blocos), então o seu modelo de 8B é o mesmo desenho com 38 blocos e 2432 canais. E o número de tokens não cresceu junto com os modelos: em 1024² ele é 4096 desde 2024. Os modelos ficaram mais largos e mais profundos; a sequência foi mantida fixa por latentes de 16 canais com \(p = 2\) (o FLUX.2 passou a um latente de 32 canais com o mesmo número de tokens).

Três propriedades vieram junto com a troca e cada uma importou mais que o número de FID:

• Uma arquitetura para tudo. O mesmo bloco, os mesmos kernels e a mesma maquinaria de treino distribuído de um LLM transferem sem mudanças: FlashAttention (um kernel de atenção fundido), paralelismo de sequência (uma sequência dividida entre GPUs), FSDP (pesos repartidos entre GPUs), checkpointing de ativações (recalcular as ativações em vez de guardá-las).
• Uma modalidade é mais tokens. Texto, quadros de vídeo, áudio, uma imagem de referência para editar: cada um é uma sequência concatenada aos tokens de imagem, não uma arquitetura nova.
• Resolução é um comprimento de sequência. Com RoPE não há compromisso arquitetural com 512 ou 1024; você muda \(N\).

Vídeo: o mesmo bloco sobre patches de espaço-tempo

Um DiT de vídeo muda a entrada, não o bloco. Um VAE 3-D comprime o clipe no tempo além do espaço — a escolha comum é 4× no tempo, 8 × 8 no espaço, 16 canais — e o patchify corta o latente em patches de espaço-tempo \(1 \times 2 \times 2\). O relatório técnico do Sora descreveu exatamente isso ("patches de espaço-tempo" entregues a um diffusion transformer), sem dar tamanhos¹¹; os modelos abertos dão:

O número de tokens é \(T \cdot \frac{H}{2} \cdot \frac{W}{2}\); o latente tem \(1 + (\text{quadros} - 1)/4\) passos de tempo (o Mochi comprime o tempo por 6 e o Movie Gen por 8, os dois com o mesmo patch \(1 \times 2 \times 2\)). Os primeiros modelos de vídeo fatoravam a atenção numa passada espacial dentro de cada quadro e numa passada temporal ao longo de cada pixel, para conter o custo. Os atuais usam todos atenção 3-D completa e o CogVideoX dá o motivo: com atenção fatorada, a cabeça de uma pessoa no quadro \(i+1\) não consegue fazer atenção direto à cabeça no quadro \(i\) — ela se moveu, então não está mais no mesmo pixel — e a consistência sofre⁹.

O mesmo desenho agora carrega outros sinais na mesma sequência: o LTX-2 gera vídeo e áudio juntos, com um fluxo de vídeo de 14B e um de áudio de 5B ligados por cross-attention¹³; o Wan 2.2 divide a remoção de ruído entre dois especialistas de 14B, um para os passos de muito ruído e outro para os de pouco ruído, com 27B no total e 14B ativos por passo²⁴.

Para onde vai o compute

O custo de um bloco Transformer do capítulo 13 vale sem mudanças:

multiplicações-adições por bloco, então a parte da atenção no trabalho é \(N/(N + 6d)\).

Um DiT de imagem não é limitado pela atenção; um DiT de vídeo é. Em 1024² as projeções e os MLPs ainda fazem 80% do trabalho. Em vídeo 720p os escores de atenção ficam com quase tudo — e a medição concorda com a fórmula: gerar um clipe de 5 segundos em 720p com o HunyuanVideo levou 945 s, 800 deles na atenção¹⁵. (A fórmula supõe um MLP de \(4d\); o do Wan é de \(2{,}7d\), o que aumenta um pouco a parte da atenção.)

Essa divisão decide quais truques de eficiência importam onde:

Para imagens, as duas primeiras linhas dominam na prática; para vídeo, a terceira é onde estão os grandes ganhos.

Onde é usado, outubro de 2026

E o que substituiu a U-Net onde: em texto para imagem, entre o SDXL (2023, U-Net) e o SD3 (2024, MMDiT); em vídeo, em 2024 — os modelos de vídeo com U-Net de 2023 (Stable Video Diffusion) não tiveram sucessores em forma de U-Net. As partes convolucionais não sumiram: o VAE que cada um desses modelos roda antes e depois do Transformer é uma CNN, assim como muitas etapas de super-resolução.

Para onde o desenho está indo

Referência de implementação

O laboratório constrói este modelo e o treina. É um DiT completo condicionado por classe, com adaLN-Zero, escrito para ser lido; os comentários dão os formatos.

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

def t_embed(t, d):                                  # t em [0, 1] -> (B, d) features senoidais
    f = torch.exp(-math.log(10000) * torch.arange(d // 2, device=t.device) / (d // 2))
    a = 1000 * t[:, None] * f[None]
    return torch.cat([a.sin(), a.cos()], dim=1)

class DiTBlock(nn.Module):
    def __init__(self, d, heads):
        super().__init__()
        self.heads = heads
        self.ln1 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)   # sem γ, β aprendidos:
        self.ln2 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)   # eles vêm de c
        self.qkv, self.proj = nn.Linear(d, 3 * d), nn.Linear(d, d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(approximate='tanh'), nn.Linear(4 * d, d))
        self.mod = nn.Linear(d, 6 * d)                                   # c -> β1 γ1 α1 β2 γ2 α2
        nn.init.zeros_(self.mod.weight); nn.init.zeros_(self.mod.bias)   # o "Zero": α = 0

    def attn(self, x):                                                   # (B, N, d) -> (B, N, d)
        B, N, d = x.shape
        q, k, v = self.qkv(x).view(B, N, 3, self.heads, d // self.heads).permute(2, 0, 3, 1, 4)
        return self.proj(F.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(B, N, d))

    def forward(self, x, c):                                             # x (B, N, d), c (B, d)
        b1, g1, a1, b2, g2, a2 = self.mod(F.silu(c))[:, None].chunk(6, dim=-1)
        x = x + a1 * self.attn(self.ln1(x) * (1 + g1) + b1)
        return x + a2 * self.mlp(self.ln2(x) * (1 + g2) + b2)

class DiT(nn.Module):
    def __init__(self, img=28, ch=1, p=4, d=96, depth=4, heads=4, n_classes=10):
        super().__init__()
        self.p, self.ch, self.g = p, ch, img // p                        # g × g patches
        self.patch = nn.Conv2d(ch, d, kernel_size=p, stride=p)           # patchify + projeção E
        self.pos = nn.Parameter(torch.zeros(1, self.g ** 2, d), requires_grad=False)  # preencher com seno-cosseno 2-D
        self.t_mlp = nn.Sequential(nn.Linear(d, d), nn.SiLU(), nn.Linear(d, d))
        self.y_emb = nn.Embedding(n_classes + 1, d)                      # último índice = "sem classe", para o CFG
        self.blocks = nn.ModuleList([DiTBlock(d, heads) for _ in range(depth)])
        self.ln = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
        self.out_mod = nn.Linear(d, 2 * d)                               # deslocamento e escala finais
        self.out = nn.Linear(d, p * p * ch)                              # token -> um patch
        for m in (self.out_mod, self.out):
            nn.init.zeros_(m.weight); nn.init.zeros_(m.bias)             # o modelo começa prevendo 0
        self.d = d

    def unpatchify(self, z):                                             # (B, N, p·p·C) -> (B, C, H, W)
        B, g, p, C = z.shape[0], self.g, self.p, self.ch
        z = z.view(B, g, g, p, p, C).permute(0, 5, 1, 3, 2, 4)           # B, C, linha, p, coluna, p
        return z.reshape(B, C, g * p, g * p)

    def forward(self, x, t, y):                                          # x (B, C, H, W), t (B,), y (B,)
        z = self.patch(x).flatten(2).transpose(1, 2) + self.pos          # (B, N, d)
        c = self.t_mlp(t_embed(t, self.d)) + self.y_emb(y)               # (B, d)
        for blk in self.blocks:
            z = blk(z, c)
        shift, scale = self.out_mod(F.silu(c))[:, None].chunk(2, dim=-1)
        z = self.ln(z) * (1 + scale) + shift
        return self.unpatchify(self.out(z))                             # velocidade, do formato de x

O treino é a loss de flow matching do capítulo 21 com \(t = 0\) como ruído: sorteie \(x_0 \sim \mathcal{N}(0, I)\), \(t \sim U(0, 1)\), forme \(x_t = (1-t)x_0 + t x_1\) e faça a regressão de model(x_t, t, y) sobre \(x_1 - x_0\), trocando o rótulo pelo índice "sem classe" em 10% das vezes, para que a mesma rede possa fazer classifier-free guidance.

Laboratório: um DiT, construído e medido

Tudo nesta página que pode ser medido em escala pequena é medido num notebook no Colab: um DiT condicionado por classe construído do zero e treinado com flow matching no Fashion-MNIST (28 × 28, dez classes de roupas). Ele precisa de GPU: no Colab, Ambiente de execução → Alterar o tipo de ambiente de execução → GPU T4 (o nível gratuito basta; a execução de referência de todas as células levou 13 minutos na GPU de um notebook Apple M2). São sete seções, cada uma terminando num checkpoint com valor de verificação:

1. patchify e unpatchify à mão: formatos, a ida e volta e a convolução com passo;
2. o bloco adaLN-Zero: identidade na inicialização (medida) e para onde vão os seus parâmetros;
3. o DiT inteiro, a sua contagem de parâmetros e a fórmula de custo contra o relógio;
4. treino com flow matching e amostragem com classifier-free guidance, julgadas por um classificador pequeno: com que frequência a amostra é da classe pedida e uma distância de Fréchet no espaço de features dele;
5. a ablação do condicionamento: tokens no contexto, adaLN e adaLN-Zero com os mesmos passos;
6. a ablação do tamanho do patch: \(p = 14, 7, 4\) — tokens, tempo por passo e qualidade;
7. o seu próprio DiT com orçamento, com uma tabela para preencher.

Abrir o laboratório no Colab — salve uma cópia no seu Drive (Arquivo → Salvar uma cópia no Drive) para rodar e editar.

Pontos principais

1. Um DiT é um ViT sobre patches latentes ruidosos: patchify, somar posições, \(L\) blocos idênticos numa única resolução e unpatchify para uma velocidade (ou ruído) do formato da entrada.
2. O tamanho do patch é um botão de compute. Reduzir \(p\) à metade quadruplica os tokens com quase o mesmo número de parâmetros e a qualidade acompanha o compute: o DiT-XL foi de FID 106 para 19,5 entre \(p = 8\) e \(p = 2\).
3. O adaLN-Zero condiciona prevendo, a partir de \(c\), a escala, o deslocamento e o portão residual da normalização de cada bloco, com o portão começando em zero para que todo bloco comece como identidade. Ele venceu os tokens no contexto e a cross-attention com o menor compute; o seu custo em parâmetros é o motivo pelo qual os modelos seguintes compartilham a modulação entre blocos.
4. O texto entra por cross-attention (texto só de leitura: Wan, PixArt), por atenção conjunta com dois fluxos (MMDiT: SD3, Qwen-Image) ou por um fluxo só (os blocos finais do FLUX, Z-Image). A atenção conjunta custa pouco compute; o que custa é o segundo conjunto de pesos.
5. As posições são RoPE por eixo — linha e coluna para imagens, mais o tempo para vídeo — e é isso que faz da resolução, da proporção e da duração uma questão de comprimento de sequência.
6. A contribuição é uma curva de escala: a qualidade acompanha os Gflops, o que justificou modelos de imagem de 12B e de 32B.
7. Imagens são limitadas pelo MLP; vídeo, pela atenção: a parte da atenção é \(N/(N + 6d)\), 20% no FLUX em 1024² e cerca de 87% em vídeo 720p — por isso é no vídeo que atenção esparsa e cache importam.