Post-Hoc Guidance for Consistency Models by Joint Flow Distribution Learning

Título: Post-Hoc Guidance for Consistency Models by Joint Flow Distribution Learning (JFDL)

Autores: Chia-Hong Hsu e Randall Balestriero (Brown University)

1. Problema Identificado

Os Modelos de Difusão (DMs) são poderosos, mas o processo de amostragem é computacionalmente caro devido à necessidade de múltiplos passos iterativos. Para controlar a fidelidade versus a diversidade na geração, utiliza-se a Guia Livre de Classificador (CFG - Classifier-Free Guidance). No entanto, a CFG em DMs exige múltiplas avaliações do modelo por amostra, aumentando o custo de inferência.

Os Modelos de Consistência (CMs) foram desenvolvidos para resolver esse problema, permitindo a geração de imagens de alta qualidade em um ou poucos passos (amostragem rápida). Contudo, as métodos existentes de guia para CMs têm limitações severas:

• Eles dependem quase exclusivamente de Distilação de Consistência (CD), que requer o conhecimento de um "professor" (um DM pré-treinado com CFG) para ensinar o CM a seguir trajetórias guiadas.
• Isso impede a aplicação de guia em modelos treinados via Consistency Training (CT), que são puramente baseados em dados e não possuem um DM professor.
• Não existe um método eficaz de guia pós-hoc (após o treinamento) para CMs que não dependa de um DM professor, criando uma lacuna significativa na aplicabilidade dos modelos CT.

2. Metodologia Proposta: JFDL

Os autores propõem o Joint Flow Distribution Learning (JFDL), um método leve e pós-hoc que habilita a guia em CMs pré-treinados sem a necessidade de um DM professor.

Conceitos Fundamentais:

• Visão de Fluxo (Flow-based): O método trata o CM pré-treinado como um solucionador de Equações Diferenciais Ordinárias (ODEs) que mapeia ruído para dados.
• Interpolação de Direções: O JFDL aprende a interpolar entre as direções de denoising de distribuições condicionadas (classe específica) e não condicionadas (ou pseudo-não condicionadas).
• Análise de Pseudo-Ruído: O núcleo teórico do método baseia-se na demonstração de que, ao combinar trajetórias condicionadas e não condicionadas, o "ruído" resultante (pseudo-ruído) segue uma distribuição Gaussiana. Isso permite que o modelo aprenda a trajetória guiada como se estivesse resolvendo um ODE com um campo de velocidade modificado.

Algoritmos:

O paper apresenta duas variações principais:

1. Naive JFDL: Requer um CM pré-treinado que já aprendeu a resolver o ODE da classe "nula" ($\emptyset$, ou seja, amostragem incondicional). O método interpola entre a solução da classe $c$ e a solução da classe $\emptyset$.
2. Random JFDL: Uma inovação crucial que remove a dependência de um solver de classe $\emptyset$. Em vez disso, utiliza uma classe aleatória ($c'$) amostrada da distribuição de dados para construir o "âncora" incondicional.
   • Descoberta Surpreendente: O método funciona eficazmente mesmo quando o CM original não foi projetado explicitamente para amostragem incondicional. A interpolação entre a classe alvo e uma classe aleatória gera um efeito de guia comparável à CFG tradicional.

Treinamento:

O JFDL é aplicado como um ajuste fino (fine-tuning) leve:

• O modelo é otimizado para minimizar uma perda combinada: a perda original de Consistency Training (ECT) + a perda JFDL.
• Utiliza-se GradNorm para balancear adaptativamente os pesos das duas perdas durante o treinamento.
• O processo é extremamente eficiente, consumindo apenas uma fração dos dados e tempo de GPU necessários para o treinamento original.

3. Contribuições Principais

1. Primeiro Método de Guia Pós-Hoc sem Professor: O JFDL é a primeira abordagem que permite adicionar capacidades de guia (estilo CFG) a CMs pré-treinados (especificamente modelos CT) sem depender de um modelo DM professor ou de distilação complexa.
2. Fundamentação Teórica em Modelos de Fluxo: O trabalho fornece uma análise teórica rigorosa, provando que a distribuição marginal do fluxo híbrido (condicionado + não condicionado) preserva a distribuição de dados e que o ruído pseudo-gerado é Gaussiano, validando a eficácia do método.
3. Validação Empírica em Dados Reais: Demonstração de que o método funciona em conjuntos de dados complexos (CIFAR-10 e ImageNet 64x64), melhorando a qualidade da geração.
4. Versatilidade (Random JFDL): A demonstração de que é possível realizar a guia sem um solver incondicional explícito, ampliando a aplicabilidade para qualquer CM pré-treinado.

4. Resultados Experimentais

Os autores avaliaram o desempenho utilizando a métrica FID (Fréchet Inception Distance) em CIFAR-10 e ImageNet 64x64.

• CIFAR-10:

  • O modelo base (ECT) sem guia apresentou um FID de 3.40 (1 passo).
  • Com Naive JFDL, o FID caiu para 3.29 (com escala de guia $\omega=1.05$).
  • Com Random JFDL, o FID caiu para 3.24 (com $\omega=1.15$).
  • Observou-se uma melhoria consistente na fidelidade para escalas de guia baixas ($\omega < 2$), replicando o comportamento da CFG em DMs.

• ImageNet 64x64:

  • O modelo base (ECT) teve FID de 5.84 (1 passo).
  • Com Naive JFDL, o FID melhorou para 4.38 ($\omega=1.80$).
  • Com Random JFDL, o FID melhorou para 4.68 ($\omega=1.80$).
  • Em amostragem de 2 passos, o JFDL também mostrou melhorias no ImageNet, embora os ganhos no CIFAR-10 tenham sido menores devido à saturação do modelo base.

• Qualidade Visual: As imagens geradas com JFDL exibem maior contraste e nitidez (fidelidade aumentada) à medida que a escala de guia ($\omega$) aumenta, com uma redução correspondente na diversidade (diversidade vs. fidelidade), exatamente como esperado em métodos de guia.

5. Significado e Impacto

Este trabalho preenche uma lacuna crítica no estado da arte de modelos generativos rápidos:

• Democratização do Controle: Permite que modelos de Consistência treinados de forma puramente orientada a dados (CT), que são mais eficientes e não dependem de modelos de difusão pesados, ganhem a capacidade de controle fino (guia) que era exclusiva dos modelos baseados em distilação.
• Eficiência Computacional: Oferece uma maneira de obter amostras guiadas de alta qualidade em um único passo, eliminando o custo de inferência de múltiplos passos dos DMs e o custo de treinamento de distilação.
• Flexibilidade: A capacidade de aplicar guia pós-hoc sem re-treinar o modelo do zero ou depender de um professor torna a tecnologia mais acessível e aplicável em cenários com recursos limitados.

Em resumo, o JFDL transforma os Modelos de Consistência em ferramentas ainda mais versáteis, combinando a velocidade de amostragem de um passo com o controle de qualidade e estilo anteriormente restrito a modelos de difusão mais lentos ou complexos.