Training-Free Refinement of Flow Matching with Divergence-based Sampling

Título: Refinamento Livre de Treinamento de Flow Matching com Amostragem Baseada em Divergência (FDS)

1. O Problema: Conflitos de Velocidade em Flow Matching

Os modelos de Flow Matching (FM) aprendem a distribuir uma distribuição alvo complexa modelando um campo de velocidade marginal. Esse campo é definido como a média das velocidades específicas de cada amostra que conectam uma distribuição prior simples (ruído) aos dados reais.

• A Limitação Fundamental: Em estados intermediários do processo de geração, as trajetórias de diferentes pares de amostras (ruído $\to$ dados) podem se cruzar. Nesses pontos de interseção, as velocidades individuais das amostras podem apontar em direções opostas (conflito direcional).
• A Consequência: Como o modelo FM aprende apenas a velocidade marginal (a média), essa média pode apontar para uma região de baixa densidade entre os modos válidos da distribuição, em vez de apontar para um modo válido.
• Resultado: Isso faz com que as trajetórias de geração se desviem para regiões de baixa densidade, resultando em imagens borradas, degradadas ou de baixa fidelidade.
• Soluções Existentes: Trabalhos anteriores tentaram resolver isso redefinindo a formulação de treinamento para modelar explicitamente a distribuição multimodal das velocidades (ex: HRF, VRFM). No entanto, essas abordagens exigem treinamento custoso do zero e não podem aproveitar modelos pré-treinados de alto desempenho já existentes.

2. Metodologia: Flow Divergence Sampler (FDS)

Os autores propõem o FDS, um framework livre de treinamento (training-free) e plug-and-play que refina os estados intermediários durante a inferência, sem modificar o modelo pré-treinado.

Ideia Central:
Em vez de tentar reestimar a distribuição de velocidades complexa, o FDS ajusta a posição do estado atual ($x_t$) para um estado vizinho ($\tilde{x}_t$) onde o mesmo campo de velocidade pré-treinado é menos ambíguo (mais confiável), antes de prosseguir com o passo do solver.

O Mecanismo Teórico (Teorema 1):
O grande insight teórico do artigo é que a discrepância entre a velocidade marginal e as velocidades individuais das amostras (que causa a degradação) pode ser quantificada e medida sem acesso aos dados de treinamento.

• Os autores provam que o erro quadrático médio residual (a medida de ambiguidade) é proporcional à divergência do campo de velocidade marginal ($\nabla \cdot u_t(x_t)$).
• Isso permite criar um surrogato livre de dados: calcular a divergência do modelo pré-treinado $u_\theta$ serve como um indicador de confiabilidade local. Uma divergência alta indica uma região de conflito/ambiguidade; uma divergência baixa indica uma região confiável.

Algoritmo de Amostragem:

1. Perturbação: Em cada passo de tempo $t$, o FDS gera um conjunto de candidatos vizinhos ao estado atual $x_t$ adicionando pequenas perturbações aleatórias (ruído gaussiano).
2. Avaliação: Para cada candidato, calcula-se a divergência do campo de velocidade do modelo.
3. Seleção: Seleciona-se o candidato com a menor divergência (menor ambiguidade).
4. Refinamento: O estado atual é atualizado para esse candidato selecionado ($\tilde{x}_t$) antes de executar o passo normal do solver (Euler, Heun, etc.) para o próximo tempo.
5. Eficiência: O método é altamente eficiente, utilizando apenas uma iteração de refinamento ($N=1$) e um único candidato ($M=1$) na maioria dos casos, aplicando-se principalmente nas fases iniciais da geração ($t < 0.5$), onde os conflitos são mais severos.

3. Contribuições Principais

1. Caracterização Teórica: Demonstração de que a discrepância entre velocidades marginais e condicionais pode ser expressa através da divergência do campo de velocidade, permitindo uma estimativa prática e livre de dados durante a inferência.
2. Framework FDS: Introdução de um método de amostragem plug-and-play que refina estados espaciais para regiões de baixa divergência, sem necessidade de retreinar o modelo.
3. Desempenho Consistente: Evidência experimental de que o FDS melhora a fidelidade da geração em diversos backbones (EDM, JiT, SiT), solvers (Euler, Heun) e tarefas, superando abordagens que apenas aumentam o número de passos de integração (NFEs) com o mesmo orçamento computacional.

4. Resultados Experimentais

Os autores validaram o FDS em várias configurações:

• Dados Sintéticos 2D: Mostraram visualmente que o FDS desvia trajetórias de regiões de alta divergência (onde o FM padrão falha) para regiões de baixa divergência, resultando em uma distribuição final mais precisa (menor distância de Wasserstein).
• Geração de Imagens (CIFAR-10 e ImageNet 256x256):
  • Em comparação com baselines ajustadas por custo computacional (aumentando os passos do solver para igualar o tempo de execução do FDS), o FDS obteve FID (Fréchet Inception Distance) significativamente menor e IS (Inception Score) mais alto.
  • Exemplo: No ImageNet com solver Heun, aumentar os passos não melhorou o FID, enquanto o FDS trouxe ganhos substanciais.
  • Qualitativamente, o FDS produziu detalhes mais finos (ex: sementes de gergolim em hambúrgueres) que os métodos padrão perdiam.
• Geração Texto-para-Imagem (SD3-Medium): Aplicado em modelos de difusão modernos, o FDS melhorou a fidelidade visual e a aderência ao texto em benchmarks como DrawBench, funcionando bem sob diferentes escalas de guia (CFG).
• Problemas Inversos (Desembaçamento e Super-resolução): O FDS foi integrado a frameworks de inferência para tarefas inversas, demonstrando melhorias consistentes em métricas como FID e LPIPS, preservando texturas e detalhes finos (ex: pelos de gatos) melhor que as bases.
• Comparação com Métodos Baseados em Treinamento: O FDS (usando o modelo EDM) superou métodos que exigem treinamento complexo para resolver cruzamentos de trajetória (como HRF), alcançando um FID muito menor com o mesmo número de parâmetros.

5. Significado e Impacto

• Eficiência e Acessibilidade: O FDS oferece uma solução prática para um problema fundamental de modelos de fluxo sem exigir retreinamento ou modificações arquiteturais. Isso permite que modelos pré-treinados de última geração (off-the-shelf) atinjam seu potencial máximo de qualidade.
• Mudança de Paradigma: Enquanto a maioria dos avanços em amostragem foca em reduzir erros de integração temporal (solvers de alta ordem), o FDS introduz uma correção espacial ativa. Ele reconhece que o problema não é apenas a precisão numérica do passo, mas a confiabilidade do campo de vetores no ponto atual.
• Generalidade: Por ser baseado apenas na divergência do campo vetorial, o método é aplicável a qualquer modelo de Flow Matching, independentemente da arquitetura ou do tipo de dado (imagem, áudio, vídeo).

Em resumo, o FDS demonstra que é possível "limpar" a trajetória de geração em tempo de inferência, guiando o modelo para regiões onde sua previsão é mais confiável, resultando em imagens de maior qualidade com custo computacional marginal.