Guiding Distribution Matching Distillation with Gradient-Based Reinforcement Learning

Título: GDMD: Guiando a Destilação de Correspondência de Distribuição com Aprendizado por Reforço Baseado em Gradientes

1. O Problema

Os modelos de difusão revolucionaram a síntese de imagens, mas sua natureza iterativa de amostragem é computacionalmente cara, impedindo interatividade em tempo real. Métodos de destilação (como a Distribution Matching Distillation - DMD) foram desenvolvidos para comprimir esses modelos em geradores de poucos passos (few-step). No entanto, a destilação tradicional frequentemente sacrifica a qualidade da imagem em prol da velocidade, pois foca estritamente em minimizar a divergência estatística (KL) em relação ao modelo professor, ignorando preferências humanas e nuances estéticas.

Recentes tentativas de integrar Aprendizado por Reforço (RL) à destilação (ex: DMDR) para guiar o modelo para regiões de alta recompensa enfrentam desafios críticos:

• Avaliação Subótima: Métodos existentes avaliam a recompensa com base nas imagens brutas geradas ($x_0$). Em estágios iniciais de treinamento, essas imagens são ruidosas e instáveis, gerando feedback de recompensa não confiável.
• Conflito de Gradientes: A otimização do RL baseada em amostras pode divergir da direção de otimização da destilação, levando a instabilidades, "hacking de recompensa" (reward hacking) e a necessidade de uma fase de "cold-start" (início frio) apenas com destilação antes de introduzir o RL.
• Limitação de Qualidade: O modelo estudante fica preso a um teto de qualidade inferior ao do professor, pois a otimização independente do RL não alinha perfeitamente com o objetivo de destilação.

2. Metodologia (GDMD)

Os autores propõem o GDMD, um novo paradigma que redefine o mecanismo de recompensa, priorizando os gradientes de destilação em vez das saídas de pixels brutos como sinal primário para otimização.

Principais Componentes:

1. Mudança de Paradigma (Gradiente vs. Amostra):

   • Em vez de avaliar a imagem gerada $x_0$, o GDMD avalia o gradiente da função de perda de DMD.
   • Os gradientes da DMD são reinterpretados como tensores de alvo implícitos ($x_{tar}$). Isso permite que modelos de recompensa existentes (como CLIP, HPS, ImageReward) avaliem diretamente a qualidade das atualizações de destilação, e não apenas o estado atual do gerador.

2. Arquitetura do Pipeline:

   • Coleta de Gradientes Consciente da Destilação (DaGC): Utiliza a aleatoriedade inerente da DMD para coletar gradientes em grupo. Emprega duas estratégias:
     • Baseada em $t$: Gera direções de gradiente diversas variando a intensidade do ruído.
     • Baseada em $\mu_{fake}$: Rastreia diferentes versões da função de pontuação (score function) durante o processo de otimização.
   • Pontuação de Gradiente Implícito (IGS): Transforma o gradiente de DMD em um alvo de perda ($x_{tar}$) que pode ser decodificado e avaliado por modelos de recompensa padrão. A equação de perda é reformulada para que a pontuação da recompensa reflita a qualidade do gradiente.
   • Otimização de Preferência Consciente de Negativos (NaPO): Inspirado no DiffusionNFT, utiliza políticas implícitas positivas e negativas para aprendizado de preferência, ajustando o campo de velocidade do modelo estudante para maximizar a recompensa sem necessidade de estimativa de verossimilhança explícita ou processos SDE complexos.

3. Vantagens Chave:

   • Estabilidade: Elimina a necessidade de "cold-start", pois a avaliação baseada em gradiente é estável desde o início do treinamento.
   • Alinhamento: Garante que a otimização do RL esteja alinhada com a direção de destilação original, evitando conflitos.
   • Eficiência: Não requer amostragem SDE (Stochastic Differential Equation) adicional para coleta de dados, utilizando apenas os gradientes já computados.

3. Contribuições Principais

• Estratégia de Pontuação Baseada em Gradientes: Pioneira na aplicação de pontuação de recompensa diretamente sobre os gradientes de otimização em vez de amostras de pixels para destilação DMD-RL.
• Pipeline Especializado: Desenvolvimento de um sistema completo incluindo DaGC, IGS e NaPO para harnessar efetivamente a orientação em nível de gradiente.
• Novo Estado da Arte (SOTA): Estabelecimento de um novo padrão para modelos de difusão de poucos passos, superando tanto métodos baseados apenas em distilação quanto os métodos híbridos anteriores (DMDR).

4. Resultados Experimentais

Os experimentos foram conduzidos em modelos baseados em SDXL e SD3, com geração em 4 passos (4-NFE).

• Desempenho Quantitativo:

  • O GDMD superou consistentemente o modelo professor (multi-passos) e todos os métodos concorrentes (LCM, Flash, DMD, DMD2, DMDR) em métricas de qualidade de imagem (Aesthetic Score, Pick Score, ImageReward) e alinhamento de texto (GenEval).
  • Em métricas não vistas durante o treinamento (Unseen Rewards), o GDMD demonstrou uma generalização superior, indicando um aprendizado de preferência robusto.
  • No benchmark GenEval, o modelo alcançou pontuações superiores em todas as dimensões (contagem, posição, atributos, cores, texto).

• Desempenho Qualitativo e Humano:

  • Redução de Artefatos: O GDMD apresentou menos distorções, saturação excessiva e alucinações de objetos em comparação com DMD e DMDR.
  • Avaliação Humana: Em um estudo de preferência humana, o GDMD venceu o modelo professor em 55,1% das comparações de qualidade de imagem e superou o DMDR em 62,7%, demonstrando uma clara preferência dos avaliadores pela qualidade e alinhamento do prompt.

5. Significância

O GDMD representa um avanço significativo na síntese de imagens eficiente. Ao resolver o conflito fundamental entre a destilação de distribuição e a otimização por reforço, o método permite que modelos de poucos passos não apenas igualem, mas superem a qualidade de seus professores de muitos passos.

Isso é crucial para aplicações em tempo real e interativas, onde a velocidade de inferência é essencial, mas a qualidade visual não pode ser comprometida. A abordagem de usar gradientes como sinal de recompensa oferece uma nova direção para a pesquisa em alinhamento de modelos generativos, sugerindo que a otimização pode ser guiada mais eficazmente através da estrutura de aprendizado do modelo do que através da avaliação de saídas brutas.