Resumo Técnico: Explorando e Explorando a Estabilidade no Flow Matching Latente

1. Declaração do Problema

O treinamento de modelos generativos, particularmente Modelos de Difusão e suas contrapartes de Flow Matching (FM), continua sendo computacionalmente caro. Esses modelos geralmente exigem conjuntos de dados massivos, recursos computacionais extensos, longas durações de treinamento e, em configurações condicionais, anotações extensas. Embora o Flow Matching ofereça uma formulação determinística para a geração de amostras via Equações Diferenciais Ordinárias (EDOs) que frequentemente reduz os passos de amostragem em comparação com Equações Diferenciais Estocásticas (EDEs), o custo de treinar esses modelos em conjuntos de dados em escala total permanece um gargalo.

A questão central abordada por este trabalho é se o tamanho do conjunto de dados e a capacidade do modelo podem ser significativamente reduzidos sem sacrificar a qualidade de geração. Especificamente, os autores investigam a estabilidade dos modelos de Latent Flow Matching (LFM) sob grandes perturbações, incluindo subamostragem de dados (poda), alterações arquitetônicas e configurações de treinamento alteradas. Eles postulam que, se as trajetórias de transporte do LFM forem inerentemente estáveis, seria possível explorar essa propriedade para derivar algoritmos de treinamento e inferência mais eficientes.

2. Metodologia

2.1. Fundamentação Teórica: Estabilidade no Flow Matching

O trabalho baseia-se na observação de que o Rectified Flow admite uma expressão de forma fechada para o campo de velocidade. Para um conjunto de dados finito, o campo de velocidade ótimo é uma soma ponderada das amostras de treinamento, onde os pesos são determinados por uma função softmax sobre as distâncias quadradas entre o estado atual e as amostras de treinamento escaladas. Os autores observam que esse softmax torna-se "acentuado" no início da fase de transporte, o que significa que uma única amostra de treinamento frequentemente domina o campo de velocidade para uma dada trajetória.

Essa percepção teórica sugere que a trajetória de transporte é determinada em grande parte pela presença de amostras específicas "dominantes". Consequentemente, remover amostras não dominantes (poda) deve deixar a maioria das trajetórias inalteradas. Os autores formalizam isso definindo uma métrica de atribuição: para uma trajetória começando em ruído $x_0$, eles determinam em qual amostra de treinamento $x_i$ ela termina. Eles hipotetizam que essa atribuição permanece invariável mesmo quando uma grande fração do conjunto de dados é removida, desde que a amostra dominante para aquela trajetória seja mantida.

2.2. Estratégias de Poda de Dados

Para aproveitar essa estabilidade, os autores propõem três critérios de poda de dados adaptados de modelos discriminativos, mas reformulados para o objetivo FM (que utiliza caminhos de ruído compartilhados e passos de tempo):

1. Pontuação baseada em Gradiente (G): As amostras são classificadas pela magnitude da norma do gradiente da perda.
2. Pontuação baseada em Perda (L): As amostras são classificadas pela magnitude da perda calculada ao longo de caminhos de ruído compartilhados.
3. Pontuação baseada em Agrupamento (C): As amostras são agrupadas em um espaço de incorporação pré-treinado (por exemplo, CLIP). Os autores propõem duas variantes específicas:
   • Proporcional ($C_p$): Seleciona amostras proporcionais ao tamanho do agrupamento.
   • Balanceada ($C_b$): Seleciona um número igual de amostras de cada agrupamento para mitigar desequilíbrios distribucionais.
   • Baseada em Kernel ($C_b^\kappa$): Seleciona amostras de forma gananciosa para corresponder à média do kernel Gaussiano do agrupamento no espaço latente.

Os autores também avaliam o "inverso" desses critérios (selecionando as pontuações mais baixas) para testar a robustez da afirmação de estabilidade.

2.3. Inferência de Grosso para Fino (C2F)

Explorando a estabilidade das trajetórias através das capacidades do modelo, os autores propõem uma estratégia de inferência em duas etapas:

1. Modelo Grosso: Uma arquitetura leve (por exemplo, DiT-S/2) é treinada em um subconjunto podado dos dados. Ele lida com a primeira fase da trajetória ($t \in [0, t_0)$).
2. Modelo Fino: Uma arquitetura de alta capacidade (por exemplo, DiT-XL/2), pré-treinada no conjunto de dados completo, lida com a segunda fase ($t \in [t_0, 1]$).

Para garantir uma transição suave no ponto de "costura" $t_0$, o Modelo Grosso é ajustado finamente com uma função de perda de continuidade. Essa perda minimiza a diferença entre os campos de velocidade dos Modelos Grosso e Fino em $t_0$. A trajetória do Modelo Fino é usada para gerar o estado alvo $x_{t_0}$ via inversão de EDO (integração reversa), garantindo que o Modelo Grosso aprenda a alinhar-se com o caminho do Modelo Fino.

3. Contribuições Principais

1. Demonstração Empírica da Estabilidade do LFM: Os autores mostram que modelos LFM treinados independentemente convergem para trajetórias de transporte altamente consistentes sob subamostragem de dados, alterações arquitetônicas e diferentes inicializações aleatórias. Essa estabilidade está ligada às propriedades teóricas do objetivo FM.
2. Algoritmos Práticos de Poda: Eles introduzem três critérios de poda (Gradiente, Perda, Agrupamento) adequados para FM. Eles demonstram que a poda agressiva (por exemplo, removendo 50-95% dos dados) não degrada o desempenho e pode até melhorar a velocidade de convergência e a qualidade de geração sob orçamentos computacionais limitados.
3. Aceleração de Inferência de Grosso para Fino: Eles propõem uma estratégia de geração de dois modelos que divide a trajetória. Essa abordagem alcança uma aceleração significativa na inferência (aproximadamente 2,15$\times$) enquanto mantém a qualidade de geração, desde que a suposição de estabilidade se mantenha entre as arquiteturas do modelo.
4. Mitigação de Viés via Poda Balanceada: Os autores mostram que a poda de agrupamento balanceada ($C_b$) pode mitigar vieses nos dados gerados (por exemplo, gênero, tom de pele) sem exigir balanceamento explícito consciente de rótulos, ao equalizar a representação de agrupamentos semânticos no conjunto de treinamento.

4. Resultados Experimentais

Os métodos foram avaliados em CelebA-HQ, FFHQ e ImageNet-1K.

• Estabilidade sob Poda: No CelebA-HQ, modelos treinados em subconjuntos disjuntos (podados em 50%) produziram imagens visualmente semelhantes quando iniciados a partir da mesma semente de ruído $x_0$. A similaridade de cosseno ArcFace entre modelos podados e não podados variou de 0,79 a 0,83, significativamente maior do que pares aleatórios (0,37).
• Métricas de Desempenho:
  • CelebA-HQ: A poda de 50% dos dados usando agrupamento balanceado ($C_b$) melhorou ligeiramente a Distância de Inception Fréchet (FID) de 24,24 (não podado) para 22,80.
  • ImageNet: A poda de 75% dos dados ($pr=0,75$) resultou em convergência mais rápida e manteve ou melhorou ligeiramente o FID em comparação com o conjunto de dados completo. Mesmo com 95% de poda, o modelo convergiu mais rapidamente inicialmente, embora o desempenho tenha degradado após um certo número de iterações.
• Robustez a Perturbações:
  • Arquitetura: O aumento da capacidade do modelo (de DiT-S/2 para DiT-XL/2) resultou em mudanças mínimas na trajetória (similaridade $s=0,81$). A troca para uma arquitetura U-Net reduziu a similaridade ($s=0,55$), mas preservou atributos grosseiros.
  • Espaço Latente: A estabilidade é sensível a mudanças na parametrização do espaço latente (por exemplo, alterar a semente do VAE ou aplicar transformações invertíveis), o que quebrou a estabilidade.
  • Objetivo: A troca de Flow Matching para difusão baseada em pontuação alterou completamente as trajetórias, indicando que a estabilidade é específica ao objetivo FM.
• Eficiência de Grosso para Fino: No CelebA-HQ, a abordagem C2F com $t_0=0,7$ reduziu o tempo de inferência de 93,95 ms/imagem (apenas Fino) para 43,53 ms/imagem (aproximadamente 2,15$\times$ de aceleração) enquanto mantinha o FID.
• Justiça: A poda balanceada ($C_b$) reduziu a divergência KL para uma distribuição uniforme em atributos (gênero, idade, tom de pele) mais efetivamente do que modelos não podados, apesar de ser agnóstica a rótulos.

5. Significado e Alegações

Os autores afirmam que a estabilidade observada é uma propriedade característica do paradigma Rectified Flow Matching sob representações latentes fixas, e não um artefato de algoritmos de treinamento específicos.

• Eficiência: O significado principal é a capacidade de reduzir o tempo de treinamento e os requisitos de dados sem sacrificar a qualidade. Os autores demonstram que, para grandes conjuntos de dados como ImageNet, a poda pode acelerar a convergência, e para conjuntos de dados menores como CelebA-HQ, pode melhorar a qualidade de geração.
• Aceleração de Inferência: A abordagem C2F oferece uma maneira principista de reduzir os custos de inferência, transferindo as etapas iniciais e grosseiras da geração para modelos leves, confiando na estabilidade da trajetória para garantir que o modelo de alta capacidade precise apenas refinar os detalhes finais.
• Controle de Viés: O trabalho sugere que a poda de conjuntos de dados, especificamente via agrupamento balanceado, pode servir como um mecanismo para controlar e mitigar desequilíbrios distribucionais em dados gerados, oferecendo um caminho para modelos generativos mais justos sem reponderação complexa consciente de rótulos.

O artigo conclui que, embora a estabilidade seja robusta a mudanças nos dados e na capacidade do modelo, ela é sensível a mudanças na estrutura do espaço latente e no próprio objetivo generativo. Essa percepção permite que profissionais projetem pipelines mais eficientes em termos de recursos, mantendo a geração de alta fidelidade.