Diffusion Transformer (DiT)

Um Transformer de difusão (DiT) é um modelo de difusão cuja rede de remoção de ruído é um transformer, em vez de uma U-Net convolucional. Divide a representação latente de uma imagem ruidosa em patches, processa-os como uma sequência de tokens com autoatenção, tal como um Vision Transformer (ViT) processa uma imagem, e prevê o ruído a remover em cada etapa de remoção de ruído. William Peebles e Saining Xie apresentaram-no em Modelos de difusão escaláveis com Transformers (ICCV 2023), onde o DiT-XL/2 atingiu uma Distância de Inception de Fréchet (FID) de 2,27 no ImageNet 256×256 com condicionamento de classe, superando o melhor resultado anterior de 3,60 obtido com difusão latente. O artigo mostrou que a qualidade da imagem melhora continuamente à medida que aumenta o poder computacional de um DiT, e desde então os denoisers Transformer substituíram as U-Nets na maioria dos grandes geradores de imagens e vídeos, incluindo Stable Diffusion 3, FLUX.1 e Sora da OpenAI.

Como funciona um Transformer de difusão#

Um DiT é executado no espaço latente do autoencoder variacional (VAE) pré-treinado usado pelo Stable Diffusion, cujo codificador reduz a resolução por um fator de 8; assim, uma imagem 256×256×3 transforma-se numa representação latente 32×32×4. Uma camada de divisão em patches separa a representação latente ruidosa em patches p×p e incorpora linearmente cada um como um token; tamanhos de patch de 2, 4 e 8 resultam em 256, 64 e 16 tokens. Embeddings posicionais seno-cosseno assinalam a origem de cada token, uma pilha de blocos Transformer processa a sequência e uma camada linear final converte cada token novamente num patch de ruído e variância previstos, que são reunidos na forma da representação latente. O passo temporal da difusão e o rótulo da classe — ou o prompt de texto, num modelo de texto para imagem — condicionam cada bloco. A geração repete esta previsão a partir de ruído puro até à última etapa, e o descodificador VAE converte a representação latente limpa em píxeis:

O artigo comparou quatro formas de introduzir a condição em cada bloco: tokens adicionais na sequência, atenção cruzada, normalização adaptativa de camadas (adaLN) e adaLN-Zero. A adaLN-Zero obteve a FID mais baixa e foi também a mais eficiente em termos computacionais. Uma pequena rede lê a soma dos embeddings do passo temporal e da classe e prevê um desvio e uma escala aplicados após cada normalização de camadas, além de uma porta que escala cada ramo residual. As portas começam em zero, pelo que cada bloco inicia o treino como uma função identidade. Este patch de código PyTorch divide uma representação latente em patches e aplica o ramo de atenção de um bloco:

import torch
from torch import nn

torch.manual_seed(0)
# A 256x256 image becomes a 4x32x32 latent; 2x2 patches turn it into 256 tokens of width 384 (DiT-S/2)
latent = torch.randn(1, 4, 32, 32)
tokens = nn.Conv2d(4, 384, kernel_size=2, stride=2)(latent).flatten(2).transpose(1, 2)

# adaLN-Zero: the timestep and class embedding regresses a shift, scale, and gate for the attention branch
condition = torch.randn(1, 384)
modulation = nn.Linear(384, 3 * 384)
nn.init.zeros_(modulation.weight)
nn.init.zeros_(modulation.bias)
shift, scale, gate = modulation(condition).unsqueeze(1).chunk(3, dim=-1)

# Self-attention branch of a DiT block, modulated by the condition and scaled by the gate
attention = nn.MultiheadAttention(384, num_heads=6, batch_first=True)
h = nn.functional.layer_norm(tokens, (384,)) * (1 + scale) + shift
out = tokens + gate * attention(h, h, h)[0]
print(f"tokens {tuple(tokens.shape)}, block is identity at initialization: {torch.equal(out, tokens)}")

Imprime tokens (1, 256, 384), block is identity at initialization: True. Os 256 tokens com largura 384 correspondem ao modelo mais pequeno, DiT-S/2, e a porta zero deixa a entrada passar sem alterações até que o treino a modifique. Um bloco DiT completo prevê um segundo desvio, uma segunda escala e uma segunda porta para o ramo MLP.

Escala com tokens e poder computacional#

O artigo treinou quatro tamanhos de modelo, do DiT-S (12 camadas, largura 384) ao DiT-XL (28 camadas, largura 1.152), cada um com tamanhos de patch 8, 4 e 2. Nos 12 modelos, a FID acompanhou o custo computacional da passagem direta em Gflops, quer o custo resultasse de uma rede mais profunda ou mais larga, quer de mais tokens. Reduzir o tamanho do patch para metade quadruplica a contagem de tokens e o custo computacional, acrescentando quase nenhuns parâmetros. O DiT-XL/2 tem 675 milhões de parâmetros, sem contar o VAE de 84 milhões de parâmetros, e usa 119 Gflops por passagem direta a 256×256. A 512×512, processa 1.024 tokens provenientes de uma representação latente 64×64×4 com 525 Gflops e atinge uma FID de 3,04. Como a autoatenção compara cada token com todos os outros, as resoluções mais elevadas e os vídeos mais longos aumentam o seu custo de forma quadrática. O Ultralytics YOLO26 evita este custo ao aplicar atenção apenas na sua etapa de escala 1/32, que tem 400 posições para uma imagem 640×640, e manter as convoluções em todas as outras partes para a deteção em tempo real.

DiT vs. U-Net e ViT#

• Denoisers U-Net: As U-Nets do DDPM e do Stable Diffusion 1.x e 2.x são codificadores-descodificadores convolucionais que trabalham em várias resoluções, com ligações de atalho entre níveis correspondentes e atenção acrescentada nas resoluções mais baixas. Um DiT mantém uma única resolução de tokens, aplica atenção global em cada bloco e tem menos escolhas de design específicas para imagens, pelo que escala da mesma forma que os Transformers nos grandes modelos de linguagem.
• Vision Transformer: O ViT codifica uma imagem para tarefas como a classificação. O DiT reutiliza a sua incorporação de patches e os seus blocos Transformer, acrescenta o condicionamento do passo temporal através de adaLN-Zero e gera uma previsão para cada patch.
• Difusão latente e correspondência de fluxo: Um modelo de difusão latente é a estrutura geral para remover ruído no espaço latente de um VAE, e a difusão ou a correspondência de fluxo é o objetivo de treino. O DiT é a rede de remoção de ruído dentro dessa estrutura, e os geradores atuais combinam backbones DiT com objetivos de correspondência de fluxo.

DiTs para texto para imagem e vídeo#

• DiT multimodal (MMDiT): O Stable Diffusion 3 processa tokens de imagem e tokens de texto com pesos separados para cada modalidade e junta-os numa única operação de atenção, permitindo que a informação flua em ambas as direções entre o prompt e a imagem. Isto melhorou a tipografia e o seguimento de prompts, e o Stable Diffusion 3.5 Large escala o design para cerca de 8 mil milhões de parâmetros.
• FLUX.1: O FLUX.1 da Black Forest Labs é um Transformer de fluxo retificado com 12 mil milhões de parâmetros para a geração de texto para imagem.
• Vídeo: O relatório técnico da Sora da OpenAI descreve um Transformer de difusão que comprime vídeos num espaço latente e divide-os em patches de espaço-tempo, para que os clipes de diferentes durações, resoluções e proporções se tornem sequências de tokens com comprimentos diferentes. Para a geração de vídeo, o maior modelo Movie Gen da Meta é um Transformer com 30 mil milhões de parâmetros, cujo contexto de 73.000 tokens de vídeo corresponde a 16 segundos de vídeo a 16 fotogramas por segundo.