WiT: Waypoint Diffusion Transformers via Trajectory Conflict Navigation

Resumo Técnico: WiT (Waypoint Diffusion Transformers)

1. O Problema: Conflito de Trajetória no Espaço de Pixels

O artigo identifica um gargalo fundamental nos modelos de geração de imagens modernos que operam diretamente no espaço de pixels (sem compressão via Autoencoders Variacionais - VAEs), como o JiT (Just-in-Time) e o Flow Matching (FM).

• Manifold Entrelaçado: Ao contrário dos espaços latentes aprendidos (que são estruturados e discriminativos), o espaço de pixels bruto é naturalmente entrelaçado e não discriminativo. Imagens visualmente semelhantes, mas semanticamente distintas, ocupam vizinhanças densas e sobrepostas no espaço de ruído.
• Conflito de Trajetória: Durante o processo de geração (denoising), os caminhos de transporte ótimo para diferentes classes semânticas tendem a convergir ou cruzar-se em regiões locais densas do espaço de ruído.
• Consequência: Para minimizar a perda de regressão (MSE), a rede neural é forçada a prever um campo de velocidade médio para essas trajetórias sobrepostas. Isso resulta em:
  • "Sangramento semântico" (semantic bleeding), onde características de classes diferentes se misturam.
  • Convergência de treinamento lenta e instável.
  • Soluções subótimas, mesmo com grandes quantidades de dados e parâmetros.

2. Metodologia: WiT (Waypoint Diffusion Transformers)

A proposta central do WiT é desemaranhar as trajetórias de geração no espaço de pixels introduzindo waypoints semânticos intermediários. Em vez de mapear diretamente do ruído para a imagem final, o processo é dividido em duas etapas estáveis e desacopladas.

Arquitetura e Componentes Principais:

1. Waypoints Semânticos (Âncoras):

   • O modelo utiliza modelos de visão pré-treinados (especificamente DINOv3) para extrair representações ricas e discriminativas das imagens.
   • Essas representações de alta dimensão são projetadas em um manifold semântico de baixa dimensão (via PCA, reduzindo para $d=64$) para criar um "waypoint" ($s_0$) que atua como uma âncora estrutural.
   • Isso transforma o problema de mapear "Ruído $\to$ Pixels" em dois problemas mais fáceis: "Ruído $\to$ Waypoint" e "Waypoint $\to$ Pixels".

2. Gerador de Waypoints Leve ($W_\psi$):

   • Um transformador pequeno e eficiente (aprox. 21M parâmetros) que, a partir do estado ruidoso atual ($z_t$) e da classe $y$, infere dinamicamente o waypoint semântico limpo ($\hat{s}_0$).
   • Ele é treinado para minimizar a perda de velocidade no espaço semântico, aprendendo a navegar do ruído até a estrutura semântica correta.

3. Gerador de Espaço de Pixels ($G_\theta$) e Just-Pixel AdaLN:

   • Um transformador principal (DiT) responsável pela síntese de alta fidelidade dos pixels.
   • Mecanismo Just-Pixel AdaLN: Diferente da condicionamento global padrão (que usa embeddings de tempo/classe uniformes), o WiT injeta o waypoint semântico previsto como uma condição espacialmente variável.
   • O waypoint é agregado ao embedding global e projetado para modular os blocos do transformador (escala, deslocamento e gate) em cada posição espacial. Isso guia a evolução dos pixels em direção ao manifold correto, evitando zonas de conflito.

Fluxo de Treinamento e Inferência:

• Treinamento em Duas Etapas: Primeiro, treina-se o Gerador de Waypoints. Em seguida, ele é congelado e usado para condicionar o Gerador de Pixels.
• Inferência: O processo começa com ruído puro. Em cada passo da ODE, o Gerador de Waypoints recalibra o waypoint semântico com base no estado atual, que então modula o Gerador de Pixels para garantir que a trajetória permaneça na direção semântica correta.

3. Contribuições Chave

1. Novo Paradigma de Geração (WiT): Propõe a primeira abordagem que resolve o conflito de trajetória no espaço de pixels puro sem recorrer a espaços latentes de VAEs (que perdem detalhes de alta frequência).
2. Mecanismo Just-Pixel AdaLN: Uma inovação arquitetônica que permite condicionamento semântico espacialmente variável, desacoplando a navegação semântica da geração de textura.
3. Fundamentação Teórica: Demonstra matematicamente (via Lei da Variância Total) que a introdução de uma condição semântica intermediária reduz a variância irreduzível do problema de transporte ótimo, encolhendo o espaço de busca e garantindo trajetórias mais estáveis.
4. Eficiência Computacional: Alcança desempenho superior com um custo computacional marginal (apenas 21M parâmetros adicionais para o gerador de waypoints).

4. Resultados Experimentais

Os experimentos foram conduzidos no dataset ImageNet 256x256.

• Desempenho Quantitativo:
  • O WiT-L/16 (Large) atingiu um FID de 2.22 e IS de 303.3 após 600 épocas.
  • Isso supera o baseline direto JiT-L/16 (FID 2.36) e rivaliza com modelos latentes pesados como DiT-XL/2 (FID 2.27), apesar de operar puramente em pixels.
  • O modelo WiT-XL/16 quebrou a barreira de FID 1.89, estabelecendo um novo estado da arte (SOTA) para modelos puramente baseados em pixels.
• Aceleração de Treinamento:
  • O WiT-L/16 alcançou o mesmo desempenho do JiT-L/16 treinado por 600 épocas em apenas 265 épocas, resultando em uma aceleração de 2.2x na convergência.
• Qualidade Visual:
  • As imagens geradas exibem maior coerência estrutural (proporções corretas em animais e cenários complexos) e preservam micro-texturas de alta frequência (penas, asas de borboleta) que costumam ser perdidas ou distorcidas em modelos baseados em VAE.

5. Significado e Impacto

O trabalho WiT é significativo porque resolve uma limitação teórica fundamental dos modelos de difusão em espaço de pixels: a ambiguidade geométrica das trajetórias de transporte ótimo.

• Sem Perda de Informação: Ao evitar VAEs, o WiT elimina o "gargalo de informação" que causa artefatos e perda de detalhes finos, oferecendo um caminho teoricamente sem perdas para a geração de imagens.
• Escalabilidade: Demonstra que é possível escalar modelos de pixels para competir com modelos latentes, desde que o problema de otimização (conflito de trajetória) seja mitigado através de ancoragem semântica inteligente.
• Futuro: Abre caminho para a geração de imagens de ultra-alta fidelidade sem a necessidade de compressão latente, unindo a estabilidade semântica de modelos de visão pré-treinados com a fidelidade bruta dos pixels.

Em resumo, o WiT prova que, ao introduzir "pontos de controle" semânticos intermediários, é possível navegar eficientemente pelo complexo espaço de pixels, resultando em modelos mais rápidos, estáveis e visualmente superiores.