The Amazing Stability of Flow Matching

Título: A Incrível Estabilidade do Flow Matching (Correspondência de Fluxo)

1. Problema e Motivação

Os modelos generativos profundos, especialmente os baseados em difusão, têm alcançado resultados notáveis na geração de amostras de alta qualidade. O Flow Matching (FM) surge como uma alternativa eficiente e simples aos modelos de difusão, utilizando um campo de velocidade $u(x, t)$ para transportar latentes ruidosos para uma variedade limpa via equações diferenciais ordinárias (ODEs).

No entanto, o treinamento desses modelos é intensivo em termos de computação e dados. Existe uma lacuna de conhecimento sobre como os modelos FM respondem a perturbações significativas nos dados de treinamento (como redução de dados) e na arquitetura. Enquanto a estabilidade de modelos de difusão sob perturbações de dados já foi observada teoricamente e empiricamente, o comportamento do Flow Matching sob essas mesmas condições permanecia pouco explorado. O artigo busca investigar se a qualidade e a diversidade das amostras geradas por FM são robustas quando submetidas a:

• Poda (pruning) agressiva de conjuntos de dados.
• Alterações na arquitetura do modelo.
• Troca de conjuntos de dados de domínio similar.

2. Metodologia

Os autores realizaram uma análise empírica da estabilidade de modelos FM treinados em subconjuntos de dados e com diferentes configurações arquiteturais.

Configuração Experimental:

• Modelo Base: Arquitetura DiT (Diffusion Transformer) adaptada para Flow Matching (FM-DiT), treinando um campo de velocidade sobre interpolantes lineares entre ruído gaussiano e dados.
• Dataset: CelebA-HQ (rostos humanos), com comparação em FFHQ.
• Métricas de Avaliação:
  • FID (Fréchet Inception Distance): Para medir a qualidade e diversidade das imagens geradas.
  • Similaridade de Cosseno ArcFace: Para quantificar a estabilidade semântica (se as imagens geradas com a mesma semente $x_0$ são visualmente semelhantes entre diferentes modelos).

Estratégias de Perturbação e Poda:
Os autores testaram três métodos de poda de dados (selecionando 50% do dataset, $p_r = 0.5$) e seus inversos (selecionando os piores 50%):

1. Aleatório (Random): Subconjunto aleatório como linha de base.
2. Baseado em Gradiente (Grad): Treina um modelo substituto pequeno para estimar a magnitude do gradiente de cada amostra. Amostras com alto gradiente são consideradas mais influentes.
3. Baseado em Perda (Loss): Seleciona amostras com base no valor de perda (loss) calculado ao longo de caminhos de ruído compartilhados.
4. Baseado em Clusters (Clust): Utiliza embeddings do CLIP e k-means para agrupar amostras semanticamente.
   • Variação Proporcional ($p$): Mantém a distribuição original dos clusters.
   • Variação Balanceada ($b$): Seleciona o mesmo número de amostras de cada cluster.
   • Critério de Seleção: Amostras mais próximas do centro do cluster (representativas) ou mais distantes (difíceis/raras).

Testes de Estabilidade:

• Subconjuntos Disjuntos: Treinar dois modelos em partes aleatórias e não sobrepostas dos dados.
• Remoção de Cluster: Treinar apenas em dados classificados como "feminino" ou "masculino" (separação binária técnica).
• Troca de Dataset: Treinar em FFHQ, mas usar o mesmo espaço latente de CelebA-HQ.
• Mudança Arquitetural: Comparar DiT-XL (grande), DiT-S (pequeno) e U-Net.

3. Principais Contribuições

1. Descoberta de Estabilidade Robusta: Demonstraram que modelos FM são surpreendentemente estáveis. Mesmo com perturbações fortes (poda de 50% dos dados, troca de dataset ou redução de capacidade do modelo), as trajetórias iniciadas com o mesmo ruído aleatório convergem para imagens visualmente e semanticamente similares.
2. Métodos de Poda Informados: Introduziram e avaliaram três métodos de poda de dados informados (Grad, Loss, Clust) para FM, analisando seu impacto qualitativo e quantitativo.
3. Melhoria via Reamostragem de Clusters: O método de reamostragem baseado em clusters que equilibra a distribuição entre diferentes grupos ($Clust_b$) conseguiu até melhorar as métricas de avaliação (FID), sugerindo que o equilíbrio dos dados é tão crucial quanto a quantidade bruta.

4. Resultados Chave

• Estabilidade Visual e Semântica:
  • Ao treinar em dois subconjuntos aleatórios disjuntos, a similaridade de cosseno ArcFace entre pares de imagens geradas pela mesma semente foi de 0.69 ± 0.12 (pares não relacionados: 0.34).
  • Ao remover um cluster inteiro (ex: apenas rostos masculinos), o modelo ainda gerou imagens similares para as classes retidas, demonstrando que a estrutura global do campo de fluxo é apenas ajustada localmente onde necessário.
  • A troca do dataset de CelebA-HQ para FFHQ manteve uma similaridade de 0.58 ± 0.15, indicando que ambos os datasets residem na mesma variedade (manifold) subjacente.
• Impacto da Arquitetura:
  • Reduzir a capacidade do modelo (DiT-XL para DiT-S) manteve alta similaridade (0.81).
  • Mudar para uma arquitetura U-Net reduziu a similaridade para 0.55, mas as atributos grosseiros ainda foram preservados, indicando estabilidade global mesmo com mudanças drásticas.
• Eficácia da Poda:
  • Grad$^{-1}$ e Loss: Selecionar amostras com gradiente ou perda mais baixos (os menos informativos) degradou significativamente o FID, mas a similaridade visual com o modelo não podado permaneceu alta (>0.79).
  • Clust$_b$: A seleção balanceada de clusters melhorou o FID (22.80 vs 24.24 do modelo original), provando que a qualidade dos dados supera a quantidade em certos cenários.
  • Perda de Informação: Mesmo métodos que pioraram o FID (como Loss) mantiveram alta similaridade semântica com o modelo completo, sugerindo que o FM é robusto a perturbações no conjunto de treinamento.

5. Significado e Conclusão

O trabalho fornece evidências empíricas de que os modelos de Flow Matching possuem uma estabilidade intrínseca notável. A conclusão central é que o campo de fluxo aprendido por esses modelos é globalmente estável; perturbações nos dados ou na arquitetura causam apenas ajustes locais no campo de velocidade, sem alterar a trajetória global que leva a imagens semanticamente consistentes.

Implicações Práticas:

• Redução de Custos: É possível treinar modelos FM de alta qualidade com menos dados (até 50% de redução ou mais), desde que a poda seja feita de forma inteligente (ex: balanceamento de clusters).
• Eficiência: A estabilidade sob mudanças arquiteturais sugere que modelos menores ou com designs diferentes podem ser usados sem perder a coerência semântica da geração.
• Generalização: A estabilidade observada sugere que a generalização do FM não depende de um único fator, mas de uma estrutura de fluxo robusta que se adapta localmente a perturbações, permitindo treinamento eficiente em grandes escalas de dados.

Em suma, o artigo demonstra que a confiabilidade dos modelos Flow Matching sob diversas perturbações é muito maior do que o esperado, oferecendo novas diretrizes para otimização de custos computacionais e seleção de dados em modelos generativos.