FAIL: Flow Matching Adversarial Imitation Learning for Image Generation

Resumo Técnico: FAIL (Flow Matching Adversarial Imitation Learning)

1. O Problema

O pós-treinamento de modelos de geração de imagens (como Flow Matching e Difusão) visa alinhar a distribuição de saída do modelo com uma distribuição alvo de alta qualidade. Atualmente, dois paradigmas dominam essa fase:

• Ajuste Fino Supervisionado (SFT): Funciona como "Clonagem Comportamental". É eficiente em dados, mas sofre do problema de desvio de distribuição (distribution shift): quando o modelo encontra estados não vistos nos dados de demonstração, a política pode desviar, levando a erros cumulativos.
• Otimização de Preferência (ex: RLHF, DPO): Aborda o desvio de distribuição, mas depende criticamente de pares de preferência (comparações humanas) ou de modelos de recompensa complexos e caros de treinar. Além disso, métodos baseados em recompensa explícita são propensos a reward hacking (o modelo otimiza a métrica em detrimento da qualidade real).

Existe uma lacuna para um método que corrija o desvio de distribuição sem depender de pares de preferência explícitos ou modelagem de recompensa complexa.

2. Metodologia

Os autores propõem o FAIL, um framework que formula o pós-treinamento de modelos de Flow Matching como um problema de Aprendizado por Imitação Adversarial (AIL). A ideia central é tratar o alinhamento como um jogo minimax entre uma política geradora ($\pi_\theta$) e um discriminador ($D_\omega$), minimizando a divergência entre a distribuição gerada e a do especialista sem recompensas explícitas.

O FAIL deriva dois algoritmos distintos para estimar gradientes, dependendo da natureza do modelo e das restrições computacionais:

• FAIL-PD (Pathwise Derivative):

  • Abordagem: "White-box". Aproveita a natureza diferenciável dos solucionadores de EDO (Equações Diferenciais Ordinárias) usados no Flow Matching.
  • Mecanismo: O gradiente do discriminador é retropropagado diretamente através dos passos de denoising até os parâmetros da política.
  • Vantagem: Fornece um sinal de aprendizado denso e de baixa variância, preservando a suavidade estrutural do manifold de fluxo. É ideal para cenários contínuos e diferenciáveis.
  • Aproximação: Utiliza uma aproximação de denoising de um único passo para reduzir o custo computacional da retropropagação em toda a trajetória.

• FAIL-PG (Policy Gradient):

  • Abordagem: "Black-box". Utiliza o Estimador de Função de Pontuação (Score Function Estimator).
  • Mecanismo: Trata a saída do discriminador como uma recompensa escalar e atualiza a política usando otimização de política (inspirado em PPO/GRPO), sem retropropagar através do gerador.
  • Vantagem: Essencial para cenários onde a retropropagação é intratável (ex: geração autoregressiva com tokens discretos) ou quando o discriminador é muito grande para calcular gradientes completos.
  • Estabilidade: Utiliza uma função de objetivo clipped e restrições de KL para evitar colapso do modelo.

Estratégias de Estabilização:
Para mitigar a instabilidade inerente ao treinamento adversarial, o FAIL emprega:

1. Imitação Híbrida: Treina com uma mistura de amostras da política online e demonstrações de especialistas (tratando estas como "rollouts perfeitos").
2. Warmup do Discriminador: Inicializa a política via Clonagem Comportamental e congela a política nos primeiros passos para garantir que o discriminador forneça um sinal de aprendizado significativo antes da adaptação adversarial.

3. Contribuições Principais

1. Formulação Teórica: Estabelece a equivalência matemática entre o pós-treinamento de Flow Matching e o Aprendizado por Imitação Adversarial, propondo o framework FAIL para preencher a lacuna entre SFT e RL.
2. Algoritmos Práticos: Desenvolve duas variantes (FAIL-PD e FAIL-PG) que cobrem desde cenários de alta performance diferenciável até cenários discretos ou de caixa-preta.
3. Regularização contra Reward Hacking: Demonstra que o FAIL pode ser integrado a métodos de RL baseados em recompensa explícita, atuando como um regularizador robusto que previne o colapso do modelo em modos de alta recompensa artificial.
4. Generalização: O framework foi validado não apenas em geração de imagens contínuas, mas também em geração de imagens discretas (LLMs autoregressivos) e geração de vídeo.

4. Resultados Experimentais

O framework foi testado ajustando o modelo FLUX.1-dev usando apenas 13.000 demonstrações (uma imagem por prompt, geradas pelo Gemini 3 Pro), simulando cenários do mundo real onde pares de preferência são inviáveis.

• Desempenho em Benchmarks:
  • UniGen-Bench: O FAIL-PD melhorou a pontuação base de 61,61 para 73,70, superando modelos de ponta como FLUX.1 Ultra e Imagen-3.0.
  • DPG-Bench: Alcançou 87,32, superando o modelo base e alcançando paridade com modelos fronteira como Seedream 3.0.
  • HPDv3 (Estética): O FAIL-PD alcançou a maior pontuação geral (11,28), superando significativamente a linha de base FLUX.1 Dev (10,43) e outros modelos competitivos.
• Comparação com RLHF e DPO:
  • Métodos de RLHF padrão mostraram instabilidade e reward hacking (melhora na recompensa simulada, mas degradação em métricas reais).
  • O FAIL superou o DPO Online, pois o discriminador dinâmico evita "falsos negativos" (rejeitar amostras de alta qualidade da própria política) e fornece feedback contínuo em vez de binário.
• Estabilidade de Treinamento:
  • O FAIL-PG converge rapidamente, mas tende a colapsar após muitos passos (450+).
  • O FAIL-PD mostra uma melhoria mais gradual, mas mantém estabilidade e melhoria monótona além de 2.000 passos, preservando a integridade do campo vetorial.
• Aplicações em Outros Modos:
  • Imagens Discretas (Xomni): O FAIL-PG melhorou significativamente a aderência ao prompt em modelos autoregressivos.
  • Vídeo (Wan2.1): O FAIL melhorou as pontuações semânticas no VBench, demonstrando generalização para síntese de vídeo.

5. Significado e Impacto

O FAIL representa um avanço significativo na otimização de modelos generativos:

• Eficiência de Dados: Prova que é possível alinhar modelos de ponta com apenas 13.000 amostras de "especialista" (uma por prompt), eliminando a necessidade de grandes conjuntos de dados de preferência humana.
• Solução para Reward Hacking: Oferece uma alternativa robusta aos métodos baseados em recompensa, que frequentemente falham ao otimizar métricas em detrimento da qualidade real.
• Versatilidade: Ao fornecer tanto uma abordagem de gradiente direto (PD) quanto de gradiente de política (PG), o FAIL torna o aprendizado por imitação adversarial aplicável a uma vasta gama de arquiteturas, desde difusão contínua até modelos de linguagem discreta e vídeo.

Em suma, o FAIL estabelece o Aprendizado por Imitação Adversarial como um paradigma versátil e eficaz para o alinhamento de modelos generativos, superando as limitações de desvio de distribuição do SFT e a complexidade de dados do RLHF/DPO.