Resumo Técnico: MixFlow Training

Problema

Este artigo aborda a discrepância entre treinamento e teste, comumente conhecida como viés de exposição (exposure bias), em modelos de difusão e flow matching. Durante o treinamento, as redes de predição recebem dados ruidosos de verdade (ground-truth) como entrada em cada passo de tempo (timestep). Por outro lado, durante o teste (amostragem), a rede recebe dados ruidosos gerados pelo passo anterior. Essa discrepância leva a erros de predição, acúmulo de erros e deriva de amostragem (sampling drift). Embora soluções anteriores tenham se concentrado em modificar o processo de amostragem (ex: Time-Shift Samplers, Epsilon Scaling) ou aplicar perturbações de entrada durante o treinamento, este trabalho propõe um novo procedimento de treinamento para mitigar diretamente esse problema.

Metodologia: MixFlow

A abordagem proposta, MixFlow, é motivada pela observação de um fenômeno de "Fluxo Lento" (Slow Flow). A análise empírica revela que os dados ruidosos de verdade mais próximos dos dados ruidosos gerados em um passo de tempo de amostragem $t$ frequentemente correspondem a um passo de tempo de maior ruído (um "passo de tempo retardado", $m_t \leq t$). Em outras palavras, os dados gerados estão atrasados em relação à trajetória da verdade fundamental (ground truth).

Para aproveitar essa observação, o MixFlow modifica o procedimento de treinamento da seguinte forma:

1. Mistura de Interpolação Retardada (Slowed Interpolation Mixture): Em vez de treinar a rede de predição $u_\theta$ em uma única interpolação $x_t$ correspondente ao passo de tempo de treinamento $t$, o MixFlow utiliza uma mistura de interpolações de um intervalo de passos de tempo "retardados".

   • Os dados ruidosos de entrada são amostrados de um conjunto $\mathcal{X}_t = \{x_{m_t} \mid m_t \in [(1-\gamma)t, t]\}$, onde $m_t$ representa um passo de tempo de maior ruído (mais lento).
   • O tamanho do intervalo é linear em relação a $t$, controlado por um coeficiente $\gamma$.

2. Função de Perda Modificada: O objetivo de treinamento otimiza a rede de predição usando a interpolação retardada $x_{m_t}$ como entrada, enquanto o alvo permanece a velocidade de verdade fundamental no passo de tempo retardado $m_t$:
   $$\mathbb{E}_{t, x_0, x_1, m_t} [\| u_\theta(x_{m_t}, t) - u^*(x_{m_t}, m_t) \|_2^2]$$
   Crucialmente, o passo de tempo de entrada passado para a rede permanece o passo de tempo de treinamento $t$, não $m_t$.

3. Estratégia de Amostragem:

   • Passo de Tempo Retardado ($m_t$): Amostrado uniformemente do intervalo $[(1-\gamma)t, t]$.
   • Passo de Tempo de Treinamento ($t$): Para garantir que os pares de $(x_{m_t}, t)$ sejam amostrados uniformemente através de todas as combinações possíveis, o passo de tempo de treinamento $t$ é amostrado de uma distribuição Beta(2, 1) ($p(t) \propto t$), em vez de uma distribuição uniforme.

4. Implementação: O método é leve, exigindo apenas modificações mínimas no código (ex: 5 linhas para RAE) em pipelines de treinamento existentes. Ele funciona como uma estratégia de pós-treinamento para modelos pré-treinados.

Principais Contribuições

• Identificação do Fluxo Lento: O artigo demonstra empiricamente que os dados ruidosos gerados no passo de tempo $t$ estão mais próximos dos dados de verdade fundamental em um passo de tempo mais lento (maior ruído), fornecendo uma base teórica para a proposta modificação de treinamento.
• Algoritmo de Treinamento MixFlow: Um novo esquema de treinamento que utiliza uma mistura de interpolações retardadas para alinhar a distribuição de treinamento mais de perto com a distribuição de teste, mitigando efetivamente o viés de exposição.
• Desempenho de Estado da Arte: A abordagem alcança melhorias significativas em várias arquiteturas de modelos (SiT, REPA, RAE) e tarefas de geração (condicional de classe e texto-para-imagem).

Resultados Experimentais

Os autores avaliaram o MixFlow na geração de imagens condicional de classe (ImageNet) e geração de texto-para-imagem (Stable Diffusion 3.5).

Geração Condicional de Classe (ImageNet)

• Modelos RAE: O MixFlow aplicado a modelos RAE alcançou novos resultados de estado da arte:
  • 256 × 256: 1.43 FID (sem orientação/guidance) e 1.10 FID (com orientação).
  • 512 × 512: 1.55 FID (sem orientação) e 1.10 FID (com orientação).
• SiT e REPA: Melhorias consistentes foram observadas em SiT (tamanhos B e XL) e modelos REPA. Por exemplo, o SiT-B/2 com orientação melhorou de 4.46 para 3.91 gFID.
• Passos de Amostragem: Os ganhos do MixFlow foram mais pronunciados com menos passos de amostragem (ex: 20 passos), sugerindo que ele aborda efetivamente os erros de aproximação inerentes à amostragem grosseira.
• Comparação: O MixFlow superou outros métodos de mitigação de viés de exposição, incluindo Input Perturbation, Epsilon Scaling e Time-Shift Samplers.

Geração Texto-para-Imagem

• Aplicado ao Stable Diffusion 3.5, o MixFlow melhorou o desempenho em todas as métricas avaliadas (GenEval, DPG-Bench, T2I-CompBench) tanto para 40 quanto para 10 passos de amostragem.
• Resultados qualitativos mostraram melhorias na precisão de contagem, relações espaciais e aderência à forma dos objetos.

Significância e Alegações

O artigo afirma que o MixFlow oferece uma solução simples e eficaz para o problema do viés de exposição ao modificar a distribuição de treinamento para considerar o fenômeno do "Fluxo Lento". Os autores enfatizam que:

1. Treinamento vs. Inferência: Ao contrário dos ajustes em tempo de inferência, o MixFlow aborda a causa raiz ao alterar o objetivo de treinamento, tornando o modelo robusto à mudança de distribuição entre treinamento e teste.
2. Compatibilidade: O método é compatível com várias arquiteturas de backbone (SiT, DiT), esquemas de interpolação (linear, preservação de variância) e paradigmas de treinamento avançados (REPA, RAE).
3. Eficiência: Alcança resultados de estado da arte com sobrecarga computacional mínima, exigindo apenas algumas linhas de modificação de código e atuando como um passo eficiente de pós-treinamento.

Os autores concluem que o MixFlow melhora significativamente a qualidade da geração, particularmente em cenários de alta resolução e amostragem de poucos passos, estabelecendo novos marcos para a geração de imagens baseada em difusão.