Posterior Augmented Flow Matching

1. Declaração do Problema: Supervisão Esparsa e Colapso do Fluxo

O artigo aborda uma limitação fundamental no Flow Matching (FM), uma classe de modelos generativos de tempo contínuo que aprendem um campo vetorial dependente do tempo para transportar amostras de uma prior simples (por exemplo, ruído gaussiano) para uma distribuição de dados complexa.

• O Problema Central: No FM padrão, o treinamento depende de supervisão esparsa, um-para-um. Para qualquer estado intermediário $z_t$ dado ao longo de uma trajetória conectando uma amostra de ruído $\epsilon$ e um ponto de dados alvo $z$, o modelo é supervisionado para prever a velocidade em direção àquele único alvo específico $z$.
• Realidade de Alta Dimensionalidade: Em espaços de alta dimensionalidade (como imagens), um único ponto intermediário $z_t$ poderia validamente residir em muitas trajetórias diferentes levando a imagens alvo diversas e válidas que satisfazem a mesma condição (por exemplo, "uma foto de um cachorro").
• Consequências:
  • Alta Variância: O gradiente em qualquer $z_t$ específico é estimado a partir de um único sinal ruidoso.
  • Colapso do Fluxo: O modelo tende a memorizar pares específicos de origem-alvo em vez de aprender o mapeamento contínuo verdadeiro. Isso leva ao "colapso do fluxo", onde entradas diversas são mapeadas para saídas excessivamente similares ou confusas, falhando em generalizar.

2. Metodologia: Posterior-Augmented Flow Matching (PAFM)

Os autores propõem o PAFM, uma generalização teoricamente fundamentada do FM que substitui a supervisão de alvo único por uma expectativa sobre a distribuição posterior de conclusões de alvos válidos.

Formulação Teórica

Em vez de minimizar a perda contra um único alvo $z$, o PAFM minimiza a perda esperada sobre um conjunto de alvos plausíveis $\{z_j\}$ extraídos da distribuição posterior $p(z_j | z_t, y)$, onde $y$ é a condição.

• Objetivo: A nova função de perda agrega gradientes de múltiplos alvos candidatos, ponderados por sua verossimilhança posterior.
• Estimador Não Viciado: Os autores provam (Teorema 1) que o PAFM é um estimador não viciado do objetivo original do FM. Em expectativa, $L_{PAFM} = L_{FM}$.
• Redução de Variância: Crucialmente, o Teorema 2 prova que, ao agregar informações de $K > 1$ trajetórias plausíveis por ponto intermediário, o PAFM limita inferiormente a variância do gradiente em comparação com o FM padrão. O fator de redução de variância está relacionado ao Tamanho de Amostra Efetivo (ESS) do esquema de amostragem.

Implementação Prática

Como a posterior verdadeira é intratável, o PAFM fatoriza o problema em duas aproximações usando Amostragem por Importância:

1. Verossimilhança do Intermediário: $p(z_t | z_j)$: Quão plausível é que o alvo $z_j$ tenha gerado o intermediário atual $z_t$? Isso é computado usando o caminho de probabilidade condicional conhecido (Gaussiano).
2. Verossimilhança da Condição: $p(y | z_j)$: Quão bem o alvo $z_j$ corresponde à condição $y$?
   • Condicionado a classe: Determinístico (1 se as etiquetas coincidirem, 0 caso contrário).
   • Condicionado a texto: Aproximado usando modelos visão-linguagem (por exemplo, CLIP/SigLIP) ou assumido uniforme se os candidatos forem pré-filtrados.

Algoritmo de Treinamento:
Para cada passo de treinamento, o modelo amostra um intermediário $z_t$. Em vez de um único alvo, ele recupera um conjunto de $K$ alvos candidatos $\{z_j\}$. O modelo prevê a velocidade em direção a uma média ponderada desses alvos, onde os pesos são proporcionais a $p(z_t | z_j) \cdot p(y | z_j)$.

3. Estratégias de Seleção de Alvos

O artigo explora várias estratégias práticas para construir o conjunto de candidatos $\{z_j\}$:

1. Vizinhos Mais Próximos no Espaço Latente (Primária): Usando FAISS para recuperar os $K$ vizinhos mais próximos da imagem de treinamento no espaço latente. Isso garante proximidade geométrica e similaridade semântica.
2. Aumentações Aleatórias: Criando alvos aplicando cortes aleatórios, redimensionamentos e rotações à imagem de origem. Isso aumenta a diversidade espacial.
3. Amostragem de Momentos VAE: Amostrando $K$ vezes da distribuição Gaussiana (momento) da codificação VAE da imagem de origem. Isso fornece pequenas perturbações em torno da codificação média.

4. Contribuições Principais

• Generalização Teórica: O PAFM é provado ser uma generalização não viciada do FM que reduz a variância do gradiente ao agregar supervisão sobre toda a posterior de alvos compatíveis.
• Mitigação do Colapso do Fluxo: Ao fornecer um sinal de supervisão mais denso, o PAFM impede que o modelo memorize pares específicos, levando a campos vetoriais mais robustos.
• Flexibilidade: O framework é agnóstico quanto à forma como os candidatos são selecionados, permitindo integração com várias modalidades de dados e arquiteturas.
• Eficiência: O método introduz sobrecarga computacional negligenciável porque os conjuntos de candidatos podem ser pré-computados (por exemplo, via índices FAISS) e carregados durante o treinamento.

5. Resultados Experimentais

Os autores avaliaram o PAFM em diferentes escalas, arquiteturas e conjuntos de dados:

• Conjuntos de Dados:
  • ImageNet-1K (Condicionado a classe): Testado nos modelos SiT-B/2 e SiT-XL/2.
  • CC12M (Texto para Imagem): Testado na arquitetura MMDiT.
• Métricas: FID50K (Distância de Inception Fréchet), sFID, Precisão, Recall.
• Ganhos de Desempenho:
  • ImageNet: O PAFM melhorou o FID em até 3,4 pontos (por exemplo, o SiT-XL/2 melhorou de 11,14 para 9,85 FID).
  • CC12M: O PAFM reduziu o FID de 10,37 (FM) para 9,45.
  • Consistência: Melhorias foram observadas em todas as escalas de modelo (SiT-B/2, SiT-XL/2) e arquiteturas (SiT, MMDiT).
• Variância do Gradiente: A análise empírica mostrou que o PAFM reduziu a variância do gradiente do mini-lote em aproximadamente 4x em comparação com o FM, levando a dinâmicas de treinamento mais estáveis.
• Custo Computacional:
  • O rendimento (throughput) diminuiu apenas 6,6%.
  • O uso de memória de pico aumentou 0,4%.
  • FLOPs por passo permaneceram inalterados.

6. Significado

O PAFM representa um passo significativo para a frente na modelagem generativa ao abordar o gargalo de "supervisão esparsa" inerente a modelos baseados em fluxo e difusão.

• Robustez: Melhora fundamentalmente a estabilidade do treinamento de modelos generativos de alta dimensionalidade, reduzindo o risco de colapso de modo e artefatos.
• Eficiência: Alcança melhorias de desempenho state-of-the-art com computação adicional mínima, tornando-o uma substituição altamente prática e direta para objetivos padrão de Flow Matching.
• Generalizabilidade: A abordagem é aplicável tanto à geração condicionada a classe quanto à geração de texto para imagem, sugerindo que poderia beneficiar uma ampla gama de modelos generativos contínuos além das arquiteturas específicas testadas.