Joint Distillation for Fast Likelihood Evaluation and Sampling in Flow-based Models

1. O Problema

Os modelos generativos modernos, como Modelos de Difusão e Flow Matching (Fluxo de Correspondência), são extremamente eficazes na geração de amostras de alta qualidade. No entanto, eles enfrentam um gargalo computacional crítico: a avaliação da verossimilhança (log-likelihood).

• Custo Computacional: Para calcular a verossimilhança exata em fluxos contínuos (CNFs), é necessário integrar numericamente a divergência do campo de velocidade ao longo de toda a trajetória da ODE (Equação Diferencial Ordinária). Isso tipicamente exige centenas a milhares de avaliações de função neural (NFEs - Neural Function Evaluations).
• O Paradoxo: Enquanto métodos recentes de distilação (como Shortcut Models, MeanFlow, Consistency Models) conseguiram reduzir drasticamente o número de passos necessários para a amostragem (gerando imagens em 1-10 passos), eles geralmente perdem a capacidade de calcular a verossimilhança.
• Consequência: Sem uma avaliação rápida e precisa da verossimilhança, técnicas fundamentais de pós-treinamento, como otimização de preferência (DPO, PPO), seleção de modelos e ajuste fino baseado em verossimilhança, tornam-se proibitivamente caras ou impossíveis para esses modelos de alta performance.

2. Metodologia: F2D2 (Fast Flow Joint Distillation)

Os autores propõem o F2D2, um framework modular que permite realizar amostragem rápida e avaliação de verossimilhança rápida simultaneamente, reduzindo o número de NFEs em duas ordens de magnitude.

Insight Central

Em Fluxos de Normalização Contínuos (CNFs), a evolução da amostra ($x_t$) e a evolução do logaritmo da densidade de probabilidade ($z_t = \log p_t(x_t)$) são governadas por um sistema acoplado de ODEs que compartilham o mesmo campo de velocidade subjacente $v_\theta$:

1. Amostragem: $\frac{dx_t}{dt} = v_\theta(x_t, t)$
2. Verossimilhança: $\frac{dz_t}{dt} = -\text{div}(v_\theta(x_t, t))$

O insight chave é que, como ambos dependem do mesmo $v_\theta$, é possível distilar conjuntamente a trajetória de amostragem e o cálculo cumulativo da divergência em um único mapa de fluxo.

Arquitetura e Treinamento

• Mapa de Fluxo Conjunto: O modelo aprende um mapa de fluxo $\Phi_\theta$ que mapeia estados de tempo $t$ para $s$ para um estado conjunto $y_t = (x_t, z_t)$.
• Parametrização: O modelo utiliza um backbone compartilhado com duas cabeças de previsão:
  1. Cabeça de Velocidade ($u_\theta$): Prevê a velocidade média para a amostragem.
  2. Cabeça de Divergência ($D_\theta$): Prevê a divergência média acumulada ao longo do intervalo.
• Objetivo de Treinamento (Loss Function): O treinamento otimiza quatro termos simultaneamente para garantir a consistência do mapa de fluxo:
  1. Perda de Amostragem (Flow Matching): Garante que a velocidade instantânea corresponda ao alvo.
  2. Perda de Consistência de Amostragem: Enforça propriedades de semigrupo ou equações de Euler/Lagrange para a trajetória $x$.
  3. Perda de Divergência Instantânea: Garante que a previsão de divergência corresponda à divergência real do campo de velocidade.
  4. Perda de Consistência de Divergência: Enforça que a divergência acumulada seja consistente com a dinâmica do sistema (usando auto-distilação).
• Estimador de Hutchinson: Para evitar o custo $O(d)$ do cálculo exato da divergência, utiliza-se o estimador de traço de Hutchinson, que fornece uma estimativa não viesada com custo $O(1)$.
• Treinamento em Etapas: O modelo de velocidade é pré-treinado primeiro (usando técnicas existentes de distilação de fluxo) para fornecer uma inicialização estável antes de adicionar a cabeça de divergência e treinar conjuntamente.

3. Contribuições Principais

1. Primeira Solução Unificada: O F2D2 é o primeiro método a permitir a avaliação precisa de log-verossimilhança em modelos baseados em difusão/fluxo com apenas alguns passos (few-step), resolvendo uma limitação de longa data.
2. Modularidade: O framework é compatível com métodos existentes de amostragem em poucos passos, como Shortcut Models e MeanFlow, exigindo apenas a adição de uma cabeça de previsão de divergência.
3. Novo Algoritmo de Guia: Os autores propõem o "Maximum Likelihood Self-Guidance". Como o modelo pode calcular a verossimilhança rapidamente, é possível otimizar o ruído inicial ($x_0$) durante a inferência para maximizar a verossimilhança, melhorando a qualidade da amostra sem necessidade de modelos de recompensa externos.

4. Resultados Experimentais

Os experimentos foram realizados nos conjuntos de dados CIFAR-10, ImageNet 64x64 e CelebA-64.

• Precisão da Verossimilhança:
  • Modelos de fluxo padrão (Flow Matching) e métodos de distilação existentes (Shortcut, MeanFlow) falham em fornecer estimativas de verossimilhança válidas em poucos passos (valores de NLL inválidos ou extremamente distantes do teacher).
  • O F2D2 produz estimativas de verossimilhança calibradas (BPD - Bits Per Dimension) muito próximas do modelo teacher (que usa 1024 passos), mesmo com apenas 1 a 8 passos.
  • Em CIFAR-10, o Shortcut-Distill-F2D2 atinge um BPD de ~3.12 (ideal) com 8 passos, comparável ao teacher.
• Qualidade da Amostra (FID):
  • O F2D2 mantém a alta qualidade de amostragem. Em alguns casos (como MeanFlow-F2D2), a qualidade da amostra (FID) até melhora em comparação com o método base, sugerindo que o sinal de verossimilhança atua como um regularizador benéfico.
• Guia de Auto-Verossimilhança:
  • Um modelo MeanFlow de 2 passos com Self-Guidance (usando F2D2) supera um modelo Flow Matching de 1024 passos em FID no CIFAR-10, demonstrando o poder do uso da verossimilhança para guiar a geração.
• Validação Sintética: Em um dataset de tabuleiro de xadrez 2D (com verossimilhança ground-truth analítica), o F2D2 recuperou a distribuição de densidade correta com apenas 1 NFE, enquanto métodos tradicionais falharam catastropicamente.

5. Significado e Impacto

O trabalho F2D2 remove uma barreira fundamental na adoção de modelos generativos de fluxo e difusão em cenários que exigem inferência probabilística rigorosa.

• Viabilidade de Otimização: Torna viáveis técnicas de fine-tuning baseadas em verossimilhança (como PPO, DPO) para modelos de difusão, que antes eram inviáveis devido ao custo computacional.
• Análise de Modelos: Permite a comparação direta e a seleção de modelos baseada em métricas de verossimilhança, algo que era difícil com modelos de poucos passos.
• Eficiência: Demonstra que é possível obter o "melhor dos dois mundos": a velocidade de inferência de modelos de poucos passos e a precisão estatística de modelos de muitos passos.

Em resumo, o F2D2 transforma modelos de fluxo de "caixas-pretas" de alta qualidade de amostragem em modelos probabilisticamente calibrados e eficientes, abrindo caminho para novas aplicações em RL, otimização e análise de dados.