Resumo Técnico: Modelagem Generativa em Um Único Passo via Fluxos de Gradiente de Wasserstein

Declaração do Problema

Modelos generativos modernos, particularmente modelos de difusão e métodos baseados em fluxo, alcançaram qualidade de amostra state-of-the-art, mas sofrem com altos custos computacionais e latência devido à sua dependência de amostragem iterativa multi-etapa. Embora existam métodos de destilação de um ou poucos passos, eles frequentemente herdam sinais de treinamento de trajetórias originalmente projetadas para geração iterativa ou dependem de dinâmicas heurísticas (por exemplo, Modelos Derivantes) que carecem de garantias de convergência principiais e são propensos ao colapso de modos. O desafio central é projetar uma dinâmica de treinamento que permita a um gerador aprender um mapeamento direto e em único passo de uma distribuição de referência para uma distribuição de dados alvo, ao mesmo tempo em que garante convergência estável e cobertura global de modos.

Metodologia: W-Flow

Os autores propõem o W-Flow, um framework que desacopla o projeto da dinâmica de treinamento da parametrização do gerador. A ideia central é prescrever a evolução da distribuição do modelo em direção à distribuição alvo usando Fluxos de Gradiente de Wasserstein (WGF) e, em seguida, treinar uma rede neural estática para comprimir essa evolução multi-etapa em um único passo de inferência.

1. Dinâmica de Treinamento via Fluxo de Gradiente de Wasserstein

Em vez de aprender uma sequência de etapas de remoção de ruído ou campos de velocidade diretamente, o W-Flow define a trajetória de treinamento como uma evolução em tempo contínuo no espaço de medidas de probabilidade.

• Funcional de Energia: A evolução é impulsionada pela descida mais íngreme de um funcional de energia $F(q) = D(q \| p)$, onde $q$ é a distribuição atual do modelo e $p$ é a distribuição de dados alvo.
• Campo de Velocidade: O campo de velocidade em tempo contínuo $V_t$ é derivado da primeira variação do funcional de energia: $V_t(x) = -\nabla \frac{\delta F}{\delta q}(q_t)(x)$.
• Discretização: O fluxo contínuo é discretizado usando um esquema de Euler explícito. Em cada etapa de treinamento $k$, a distribuição do modelo $q^{(k)}$ é atualizada via um mapa de transporte local $g^{(k)}(x) = x + \eta V^{(k)}(x)$, onde $\eta$ é o tamanho do passo.

2. Instanciação com Divergência de Sinkhorn

Para garantir que o campo de velocidade seja computacionalmente tratável e geometricamente bem-comportado em altas dimensões, os autores instanciam o funcional de energia usando a divergência de Sinkhorn ($S_\epsilon$), uma métrica de transporte ótimo (OT) regularizada entropicamente.

• Formulação do Campo de Velocidade: O campo de velocidade induzido é dado pela diferença entre duas projeções baricêntricas:
  $$V^\epsilon_{q,p}(x) = T^\epsilon_{q,p}(x) - T^\epsilon_{q,q}(x)$$
  Aqui, $T^\epsilon_{q,p}$ transporta massa da distribuição gerada $q$ para o alvo $p$, enquanto $T^\epsilon_{q,q}$ representa um termo de auto-transporte que impõe consistência com a distribuição atual.
• Coordenação Global: Diferentemente de interações locais de kernel (por exemplo, em MMD ou Modelos Derivantes), o campo de velocidade de Sinkhorn é derivado de um plano de OT global, garantindo que os movimentos das partículas sejam coordenados em todo o lote para prevenir o colapso de modos.
• Estimador de Dois Lotes: Para evitar artefatos de auto-matching onde partículas correspondem a si mesmas com custo zero, os autores empregam uma estratégia de dois lotes. Eles calculam o termo de auto-transporte $T^\epsilon_{q,q}$ usando um segundo lote independente de amostras da distribuição atual, em vez de mascarar a diagonal da matriz de custo de um único lote.

3. Treinamento do Gerador

O gerador $f_\theta$ é treinado para imitar as atualizações discretas de partículas definidas pelo WGF.

• Objetivo: A função de perda minimiza a distância quadrada entre a saída do gerador e a posição de partícula "alvo" derivada da atualização de Sinkhorn:
  $$\mathcal{L}(\theta) = \mathbb{E}_{z \sim p_{ref}} \| f_\theta(z) - \text{sg}(f_\theta(z) + \eta V^\epsilon(\hat{q}_\theta, \hat{p})) \|^2$$
  onde $\text{sg}(\cdot)$ denota o operador stop-gradient, tratando a posição de partícula atualizada como um alvo de regressão fixo.
• Guia Livre de Classificador (CFG): O framework suporta CFG modificando diretamente o campo de velocidade. Em vez de misturar distribuições linearmente, os autores propõem guia de velocidade, que adiciona um termo proporcional à diferença entre campos de velocidade condicionais e incondicionais. Esta abordagem mostra-se capaz de recuperar a distribuição alvo exponencialmente inclinada padrão ao usar divergência KL e fornece uma extensão geometricamente consistente para a divergência de Sinkhorn.

Contribuições Principais

1. Framework Principial de Um Único Passo: O W-Flow estabelece um novo paradigma para modelagem generativa em um único passo, fundamentando a dinâmica de treinamento em fluxos de gradiente de Wasserstein, oferecendo uma interpretação variacional clara ausente em abordagens heurísticas como Modelos Derivantes.
2. Dinâmica Baseada em Sinkhorn: A instanciação do WGF com divergência de Sinkhorn fornece uma regra de atualização eficiente baseada em transporte ótimo que captura discrepâncias distribucionais globais e melhora a cobertura de modos em comparação com métodos baseados em kernel ou de interação local.
3. Convergência Teórica: Os autores provam que a dinâmica de partículas de amostra finita do W-Flow converge para a dinâmica distribucional em tempo contínuo sob suposições adequadas de regularidade, fornecendo uma justificação teórica para o processo de aprendizado.
4. Estimador de Dois Lotes: Uma nova estratégia de dois lotes para estimar termos de auto-transporte é introduzida, eliminando o viés de auto-matching inerente a estimadores de lote único e heurísticas de mascaramento de diagonal.

Resultados Experimentais

Os autores avaliam o W-Flow em ImageNet 256×256 e FFHQ.

• Desempenho State-of-the-Art: O W-Flow alcança um novo FID state-of-the-art de 1,29 (escala XL) e 1,35 (escala L) para geração condicional a classes em um único passo no ImageNet, superando significativamente métodos anteriores de um único passo (por exemplo, Modelos Derivantes, MeanFlow) e superando muitos modelos de difusão multi-etapa (por exemplo, LightningDiT, SiT), apesar de usar apenas uma avaliação de função (NFE).
• Eficiência de Amostragem: Em comparação com modelos de difusão multi-etapa com pontuações FID semelhantes, o W-Flow produz uma amostragem aproximadamente 100× mais rápida (por exemplo, 77,88 imagens/segundo vs. 0,93 imagens/segundo para SiT-XL/2 em uma única GPU H100).
• Cobertura de Modos: Em experimentos com distribuições desbalanceadas (por exemplo, 95% de rostos de idosos, 5% de rostos de crianças), o W-Flow captura com sucesso modos minoritários, enquanto Modelos Derivantes exibem colapso de modos, ignorando completamente a classe minoritária.
• Velocidade de Convergência: O W-Flow converge substancialmente mais rápido durante o treinamento, alcançando o desempenho de um Modelo Derivante treinado por 1280 épocas em apenas 384 épocas.
• Transferência de Domínio: O método demonstra forte capacidade em tarefas de transferência de domínio (por exemplo, tradução de idade no FFHQ), produzindo trajetórias de transporte mais curtas e coerentes em comparação com bases heurísticas.

Significado e Alegações

O artigo afirma que fluxos de gradiente de Wasserstein fornecem uma base principial e eficaz para modelagem generativa rápida e de alta fidelidade. Ao separar o projeto da dinâmica de treinamento da parametrização do gerador, o W-Flow demonstra que um gerador de único passo pode ser treinado para alcançar desempenho comparável a processos iterativos multi-etapa complexos. Os autores enfatizam que a coordenação global fornecida pela dinâmica de transporte ótimo é fundamental para estabilizar o treinamento e garantir cobertura abrangente de modos, abordando limitações críticas das abordagens existentes de um único passo. Os resultados sugerem que o W-Flow pode servir como uma alternativa escalável para aplicações que exigem geração de alta qualidade e baixa latência, como ferramentas criativas em tempo real e simulação.