Resumo Técnico: Drift Flow Matching (DFM)

Declaração do Problema

Modelos generativos iterativos, como Flow Matching e modelos de Difusão, demonstraram fortes capacidades de "escalabilidade em tempo de teste", onde o aumento do número de passos de inferência melhora a qualidade da geração, a controlabilidade e a cobertura da distribuição. No entanto, esses métodos exigem múltiplos passos de inferência, o que pode ser computacionalmente caro. Por outro lado, Modelos de Deriva (Drift Models) oferecem geração eficiente em um único passo, aprendendo um mapa de empurrão direto de uma distribuição fonte para uma distribuição alvo. Ainda assim, seu paradigma de geração direta carece de um mecanismo natural para melhorar os resultados aumentando os passos de inferência em tempo de teste, limitando sua flexibilidade e adaptabilidade a diferentes requisitos de qualidade e eficiência.

Metodologia: Drift Flow Matching (DFM)

Os autores propõem o Drift Flow Matching (DFM), um framework que unifica a eficiência dos Modelos de Deriva de um único passo com a escalabilidade do Flow Matching de múltiplos passos. O DFM alcança isso aprendendo um modelo de transporte de dois tempos que pode mover amostras de qualquer tempo atual $t$ para qualquer tempo futuro $r > t$ ao longo de um caminho de probabilidade do Flow Matching.

Mecanismos Principais

1. Construção de Marginais de Dois Tempos:
   Diferentemente do Flow Matching padrão, que regressa um alvo de velocidade condicional pontual, o DFM opera em pares de distribuições marginais $(p_t, p_r)$. Ele define um caminho condicional $X^Z_t = \alpha(t)X^Z_0 + \beta(t)X^Z_1$ e $X^Z_r = \alpha(r)X^Z_0 + \beta(r)X^Z_1$ baseado em um acoplamento independente das distribuições fonte ($X_0$) e alvo ($X_1$). Essa construção induz distribuições marginais $p_t$ e $p_r$ em passos de tempo arbitrários.

2. Parametrização de Velocidade Média:
   O DFM parametriza um campo de velocidade média $u_\theta(x_t, t, r)$ para transportar a distribuição atual $p_t$ para a distribuição futura $p_r$. O mapa de transporte é definido como:
   $$T^\theta_{t,r}(x_t) = x_t + (r - t)u_\theta(x_t, t, r)$$
   Essa formulação permite que o modelo lide com tamanhos de passo arbitrários $(r-t)$. No limite quando $r \to t$, a velocidade média recupera a velocidade marginal instantânea do Flow Matching padrão.

3. Supervisão Baseada em Deriva:
   Em vez de regressar um vetor de velocidade, o DFM emprega um objetivo baseado em deriva para transportar a distribuição prevista $q^\theta_{t,r}$ (gerada por $T^\theta_{t,r}$) em direção à marginal alvo verdadeira $p_r$.

   • Campo de Deriva: Um campo de deriva $V_{q^\theta_{t,r}, p_r}$ é construído usando uma atração ponderada por kernel em direção às amostras alvo ($V^+_{p_r}$) e um termo de auto-correção repulsivo para longe da distribuição atual do modelo ($V^-_{q^\theta_{t,r}}$).
   • Função de Perda: O objetivo de treinamento minimiza a magnitude do campo de deriva usando uma formulação de gradiente interrompido (stop-gradient):
     $$L_{DFM}(\theta) = \frac{1}{2} \mathbb{E}_{(t,r), X_t \sim p_t} \left[ \| \hat{X}^\theta_r - \text{sg}(\hat{X}^\theta_r + V_{q^\theta_{t,r}, p_r}(\hat{X}^\theta_r)) \|^2 \right]$$
   • Estimativa Agrupada: Crucialmente, o campo de deriva e a perda devem ser computados agrupadamente para cada par de tempos amostrado $(t, r)$. Misturar amostras de diferentes pares de tempos corromperia a estimativa da deriva, pois as marginais alvo $p_r$ diferem para cada par.

Contribuições Principais

• Framework Unificado: O DFM preenche a lacuna entre Modelos de Deriva eficientes de um único passo e métodos iterativos de Flow Matching, criando um paradigma único que suporta tanto a geração direta quanto o refinamento iterativo.
• Escalabilidade em Tempo de Teste: O framework desbloqueia capacidades de escalabilidade em tempo de teste para métodos baseados em deriva. Os usuários podem trocar velocidade de amostragem por qualidade de geração ajustando o número de passos de inferência (NFE) sem retreinamento.
• Princípio de Aprendizado Inovador: O artigo introduz um princípio de aprendizado que migra da regressão de velocidade pontual (padrão no Flow Matching) para o transporte de distribuição marginal via supervisão por deriva. Isso evita a instabilidade frequentemente associada a métodos de auto-distilação usados em outras abordagens de geração em um único passo.
• Garantias Teóricas: Os autores fornecem análise teórica mostrando que o DFM decompõe a geração de pontos finais em problemas de transporte marginal de curto alcance controlados pela distância de Wasserstein ($W_2$) e que o método recupera o Flow Matching padrão no limite de passo infinitesimal.

Resultados Experimentais

Experimentos extensivos em dados sintéticos, MNIST, FFHQ, ImageNet-1k e tarefas de controle robótico demonstram a eficácia do DFM:

• Dados Sintéticos: Visualizações mostram que o DFM pode gerar amostras de alta qualidade em um único passo, enquanto melhora ainda mais a cobertura da distribuição e a qualidade das amostras à medida que o NFE aumenta.
• Geração Condicional (MNIST e FFHQ): O DFM alcança desempenho competitivo com inferência em um único passo (igualando as linhas de base de Modelos de Deriva) e melhora consistentemente as métricas (EMD, FID) à medida que o NFE aumenta para 2, 5 ou 10.
• Síntese de Imagem em Grande Escala (ImageNet-1k): No conjunto de dados ImageNet-256x256, o DFM alcança um FID competitivo de 1,52 com NFE=1. Aumentar o NFE para 10 reduz ainda mais o FID para 1,31, demonstrando escalabilidade eficaz em tempo de teste comparável a métodos iterativos de Flow Matching.
• Controle Robótico: Em tarefas de manipulação robótica (por exemplo, Lift, Can, PushT), o DFM preserva a capacidade de inferência em um único passo da Política de Deriva (Drift Policy), enquanto exibe ganhos claros de desempenho com o aumento dos passos de inferência, superando tanto a Política de Deriva de um único passo quanto as linhas de base de Políticas de Difusão de múltiplos passos em tarefas específicas.

Significado e Alegações

O artigo alega que o DFM estabelece um paradigma generativo mais flexível, onde a geração direta e o refinamento iterativo não são mais escolhas de design mutuamente exclusivas, mas podem ser vistas dentro de um único framework. Ao conectar a modelagem generativa de deriva com a geração iterativa baseada em fluxo, o DFM permite que os processos de geração adaptem o cálculo de amostragem a diferentes requisitos de qualidade e eficiência. Os autores enfatizam que essa abordagem fornece uma maneira inovadora de alcançar forte geração em um único passo, mantendo os benefícios de escalabilidade dos métodos iterativos, abordando uma limitação chave dos Modelos de Deriva anteriores.