Fine-Tuning Diffusion Models via Intermediate Distribution Shaping

Este artigo propõe o P-GRAFT, um método que otimiza o ajuste fino de modelos de difusão ao moldar distribuições em níveis de ruído intermediários para melhorar a maximização de recompensas e corrigir erros de aprendizado sem recompensas explícitas, demonstrando ganhos significativos em benchmarks de geração de imagens e moléculas.

O essencial

Imagine que você tem um artista muito talentoso, mas um pouco "preguiçoso" ou que aprendeu algumas coisas erradas na escola. Esse artista é um Modelo de Difusão (como o Stable Diffusion), que cria imagens, músicas ou moléculas a partir do nada (ruído).

O problema é que, às vezes, esse artista não segue exatamente o que você pede, ou cria coisas que não são perfeitas. A ideia deste artigo é: "Como corrigir esse artista sem ter que ensiná-lo tudo do zero?"

Os autores propõem duas técnicas principais, que vamos explicar com analogias do dia a dia:

1. A Técnica do "Rascunho Intermediário" (P-GRAFT)

O Problema:
Normalmente, para treinar um artista, você pede para ele fazer 100 desenhos, olha o resultado final, e diz: "Gostei deste, não gostei daquele". Mas, no mundo da Inteligência Artificial, olhar apenas o resultado final é como tentar adivinhar onde um carro vai bater olhando apenas a foto do acidente. É difícil saber exatamente qual movimento do motorista causou o erro. Além disso, o "ruído" inicial é tão bagunçado que é difícil aprender com ele.

A Solução (P-GRAFT):
Em vez de esperar o desenho ficar 100% pronto para julgar, os autores sugerem julgar o artista no meio do caminho.

• A Analogia: Imagine que você está ensinando alguém a cozinhar um bolo.
  • Método antigo: Você deixa a pessoa fazer o bolo inteiro, assa, e só depois diz: "Esse bolo ficou ruim". A pessoa não sabe se o erro foi na farinha, no tempo do forno ou no açúcar.
  • Método P-GRAFT: Você deixa a pessoa misturar os ingredientes e, no meio do processo, você olha a massa e diz: "Essa massa está ótima, vamos continuar com ela". Se a massa estiver ruim, você descarta e pede para tentar de novo.

Por que funciona?
O artigo explica que, no meio do processo (quando a imagem ainda é um "rascunho" meio borrado), é mais fácil para o modelo aprender a corrigir o erro do que quando a imagem já está quase pronta. É como corrigir a direção de um carro quando ele está saindo da garagem (fácil) em vez de tentar corrigir quando ele já está a 100km/h (difícil e perigoso).

Resultado: Eles testaram isso em Stable Diffusion e o modelo ficou muito melhor em seguir instruções (como "um gato vermelho em cima de um cachorro azul") do que os métodos antigos.

---

2. A Técnica do "Ruído Inverso" (Inverse Noise Correction)

O Problema:
Existem modelos que não usam "ruído" aleatório, mas sim um fluxo contínuo (chamados Flow Models). Imagine que esses modelos são como um rio que leva água (imagens) de uma montanha (ruído) até o vale (imagem final). O problema é que, às vezes, a água chega suja ou com detritos, mesmo que o rio esteja bem desenhado. Isso acontece porque o mapa do rio (o modelo treinado) tem pequenos erros.

A Solução (Inverse Noise Correction):
Em vez de tentar consertar o rio inteiro (o que é caro e difícil), eles propõem consertar apenas a fonte da água.

• A Analogia: Imagine que você tem um filtro de água defeituoso que deixa a água sair com gosto estranho.
  • Método antigo: Tentar limpar a água em cada torneira da casa (gasta muita energia).
  • Método Inverso: Eles pegam a água que já saiu (a imagem final), "desfazem" o caminho (como se a água voltasse a subir a montanha) e veem de onde ela saiu. Descobrem que a fonte estava um pouco suja. Então, eles criam um novo filtro (o "Corretor de Ruído") que limpa a água antes de ela entrar no rio principal.

Por que funciona?
É muito mais fácil e barato limpar a água na fonte do que tentar limpar cada gota que sai da torneira. Com esse "filtro extra" pequeno e leve, eles conseguem melhorar a qualidade das imagens geradas, gastando menos energia de computador.

---

Resumo dos Resultados

Os autores testaram essas ideias em várias áreas:

1. Gerar Imagens: O modelo aprendeu a seguir prompts (instruções) muito melhor do que os concorrentes.
2. Gerar Moléculas: Conseguiram criar mais moléculas estáveis (úteis para remédios) sem gastar tanto tempo.
3. Layouts: Melhoraram a organização de elementos em páginas (como jornais ou sites).

A Lição Principal:
Não precisa ser um gênio para consertar um sistema complexo. Às vezes, basta olhar para o meio do caminho (para ajustar o aprendizado) ou limpar a fonte (para corrigir o erro inicial). O artigo mostra que essas "truques" simples, mas inteligentes, podem fazer a Inteligência Artificial criar coisas muito melhores, mais rápido e gastando menos energia.

Resumo técnico

1. Problema e Motivação

Os modelos de difusão pré-treinados são amplamente utilizados para tarefas generativas, mas frequentemente necessitam de ajuste fino (fine-tuning) para corrigir erros de aprendizado ou alinhar-se a aplicações específicas (como geração de imagens baseada em texto).

O desafio central abordado no artigo é a dificuldade de aplicar métodos de Otimização de Política Próxima (PPO) e regularização por Divergência KL em modelos de difusão. Diferentemente dos modelos autoregressivos (como LLMs), a verossimilhança marginal em modelos de difusão é intratável, tornando a implementação direta da regularização KL impossível. As abordagens existentes, como a regularização KL de trajetória, muitas vezes resultam em treinamento instável ou desempenho subpar.

Além disso, métodos baseados em gradiente de política (como DDPO) enfrentam problemas de viés inicial e alta variância quando escalados para grandes conjuntos de dados. O artigo busca uma alternativa que permita a maximização de recompensas com regularização KL marginal sem calcular a verossimilhança marginal explicitamente.

2. Metodologia

Os autores propõem uma abordagem baseada em Amostragem por Rejeição para o ajuste fino, unificando conceitos teóricos e práticos em três pilares principais:

A. GRAFT (Generalized Rejection sAmpling Fine-Tuning)

O trabalho começa por unificar variantes existentes de ajuste fino por rejeição (como RAFT e RSO) sob um framework chamado GRAFT.

• Conceito: O GRAFT demonstra que a amostragem por rejeição (selecionando um subconjunto de amostras com base em recompensas) realiza implicitamente a maximização de recompensas regularizada por KL.
• Mecanismo: Ao rejeitar amostras de baixa recompensa e treinar o modelo apenas nas amostras aceitas, o modelo aprende a amostrar de uma distribuição "inclinada" (tilted distribution) que maximiza a recompensa esperada menos a divergência KL em relação ao modelo de referência.
• Vantagem: Isso contorna a intratabilidade da verossimilhança marginal, permitindo a regularização KL efetiva em modelos de difusão.

B. P-GRAFT (Partial-GRAFT) e o Trade-off Viés-Variância

A principal contribuição metodológica é o P-GRAFT, que introduz o conceito de Moldagem de Distribuição Intermediária.

• Ideia Central: Em vez de ajustar o modelo apenas no estado final (imagem limpa), o P-GRAFT realiza o ajuste fino apenas até um passo de desruído intermediário ($t$). As recompensas são calculadas com base na imagem final gerada, mas o treinamento ocorre nos vetores latentes intermediários.
• Justificativa Teórica (Viés-Variância):
  • Variância: A recompensa condicional ao estado intermediário $X_t$ tem variância alta quando $t$ é próximo ao ruído puro (início do processo), tornando o sinal de recompensa ruidoso.
  • Viés: No entanto, a função de pontuação (score function) do modelo torna-se exponencialmente mais simples à medida que $t$ se aproxima do ruído puro (convergindo para a pontuação de uma distribuição Gaussiana).
  • Otimização: O P-GRAFT encontra um ponto intermediário ideal onde a complexidade de aprendizado (viés) é reduzida o suficiente para superar o ruído da recompensa (variância), facilitando o treinamento.
• Implementação: O modelo é ajustado apenas para os passos de desruído de $T$ até $t_{intermediário}$. Para os passos restantes (até a imagem final), utiliza-se o modelo de referência (pré-treinado).

C. Correção de Ruído Inverso (Inverse Noise Correction) para Modelos de Fluxo

Para modelos de fluxo (Flow Models), que são determinísticos e baseados em ODEs, os autores propõem uma técnica para corrigir erros de aprendizado sem recompensas explícitas.

• Problema: Erros no modelo pré-treinado levam a uma distribuição final que não corresponde perfeitamente aos dados reais.
• Solução: Aproveitando a reversibilidade exata dos modelos de fluxo, o método "inverte" o processo:
  1. Gera-se amostras do modelo pré-treinado.
  2. Inverte-se a ODE (usando Euler Reverso) para mapear essas imagens de volta para o espaço de ruído inicial.
  3. A distribuição resultante no espaço de ruído é chamada de Distribuição de Ruído Inverso.
  4. Treina-se um pequeno adaptador ("Noise Corrector") para aprender a mapear ruído Gaussiano padrão para essa distribuição de ruído inverso corrigida.
• Resultado: Ao gerar novas imagens, o sistema primeiro corrige o ruído inicial e depois aplica o modelo de fluxo original, corrigindo os erros de distribuição sem necessidade de recompensas externas.

3. Contribuições Principais

1. Unificação Teórica (GRAFT): Estabelece que métodos baseados em rejeição são equivalentes à maximização de recompensas regularizada por KL, resolvendo o problema da intratabilidade da verossimilhança marginal em difusão.
2. Moldagem Intermediária (P-GRAFT): Propõe um novo paradigma de ajuste fino que treina apenas a parte inicial do processo de desruído, provando teoricamente e empiricamente que isso melhora a estabilidade e a qualidade devido a um melhor trade-off viés-variância.
3. Correção de Ruído Inverso: Introduz um método eficiente para melhorar modelos de fluxo pré-treinados corrigindo a distribuição de ruído inicial, eliminando a necessidade de recompensas explícitas.
4. Generalização: O framework é aplicado com sucesso em modelos contínuos (difusão), modelos de fluxo e modelos de difusão discreto-contínua (IGD).

4. Resultados Experimentais

Os métodos foram avaliados em várias tarefas: geração de imagem texto-para-imagem (T2I), geração de layout, geração de moléculas e geração de imagem incondicional.

• Geração Texto-para-Imagem (Stable Diffusion v2):

  • O P-GRAFT superou consistentemente métodos baseados em gradiente de política (como DDPO) e o modelo base (SDv2) em benchmarks populares (GenAI-Bench, T2ICompBench++, GenEval).
  • Em termos de VQAScore, o P-GRAFT mostrou uma melhoria relativa de 8,81% sobre o modelo base.
  • O P-GRAFT com passo intermediário em $0.25N$ (25% do tempo total de desruído) obteve os melhores resultados, validando a teoria do trade-off viés-variância.
  • O método superou o DDPO mesmo com menos chamadas de gradiente e amostras de treinamento.

• Geração de Layout e Moléculas:

  • Em geração de layout (PubLayNet), o P-GRAFT melhorou tanto a métrica de alinhamento quanto o FID, mantendo-se mais próximo da distribuição de pré-treinamento que o GRAFT padrão.
  • Em geração de moléculas (QM9), o P-GRAFT alcançou maior estabilidade molecular (92,61% vs 90,50% do baseline) com apenas 1x o número de rodadas de amostragem necessárias para o GRAFT (que precisou de 9x), demonstrando eficiência de amostragem.

• Correção de Ruído Inverso (Modelos de Fluxo):

  • Na geração de imagens incondicionais (CelebA-HQ e LSUN-Church), a correção de ruído inverso melhorou significativamente o FID (Fréchet Inception Distance).
  • Além da qualidade, o método permitiu geração mais rápida: um pipeline com 100 passos do corretor + 100 passos do modelo pré-treinado superou o modelo pré-treinado com 1000 passos, com custos computacionais (FLOPs) reduzidos.

5. Significância e Impacto

Este trabalho é significativo por várias razões:

• Superação de Limitações do PPO: Oferece uma alternativa robusta e estável ao PPO para ajuste fino de modelos de difusão, evitando os problemas de instabilidade e viés inicial associados a métodos de gradiente de política em larga escala.
• Eficiência Computacional: O P-GRAFT e a Correção de Ruído Inverso demonstram que é possível obter ganhos de qualidade significativos com menos recursos computacionais (menos gradientes, menos passos de inferência).
• Fundamentação Teórica Sólida: A conexão entre amostragem por rejeição e regularização KL, bem como a análise de viés-variância para distribuição intermediária, fornece uma base teórica rigorosa para futuras pesquisas em ajuste fino de modelos generativos.
• Versatilidade: A metodologia não se limita a imagens; ela se aplica a domínios complexos como química (moléculas) e layout, sugerindo um caminho promissor para o alinhamento de modelos generativos em diversas áreas.

Em resumo, o artigo propõe uma mudança de paradigma no ajuste fino de modelos de difusão, focando na moldagem da distribuição em estágios intermediários e na correção da distribuição de ruído inicial, resultando em modelos mais alinhados, estáveis e eficientes.