Resumo Técnico: Ajuste Fino de Difusão com Moment-Matching Distillation Recompensado

Declaração do Problema

Os modelos de difusão tornaram-se o padrão para síntese de imagens de alta fidelidade e previsão científica devido aos seus objetivos de treinamento estáveis. No entanto, a implantação prática é dificultada pelo alto custo computacional da inferência, que requer dezenas a centenas de etapas de denoising (redução de ruído). Embora técnicas de destilação comprimam esses modelos em geradores de poucos passos, e o ajuste fino por Aprendizado por Reforço (RL) os alinhe com funções de recompensa específicas (ex: preferências humanas ou métricas científicas), combinar essas duas fases permanece um desafio.

Abordagens existentes para fundir destilação e otimização de recompensa enfrentam limitações significativas:

1. Fragilidade no Treinamento Conjunto: Mesclar destilação e ajuste fino de recompensa em uma única fase frequentemente causa o desvio dos exemplos gerados para fora da distribuição de entrada do modelo professor, invalidando o sinal de destilação.
2. Problemas de Memória e Viés: O ajuste fino de um modelo já destilado via backpropagation de múltiplos passos (ex: DRaFT) é intensivo em termos de memória. Truncar o backpropagation (ex: ReFL) introduz viés ao avaliar recompensas em latentes ruidosos.
3. Restrições Estruturais: Métodos como HyperNoise evitam o backpropagation em cadeia perturbando o ruído inicial, mas são estruturalmente limitados a mudanças de baixa frequência na imagem, limitando sua capacidade de otimizar recompensas estruturais complexas de alta frequência.

Metodologia: Rewarded Moment Matching Distillation (RMMD)

Os autores propõem o RMMD, um framework de duas fases que conecta de forma principista a destilação e o ajuste fino de recompensa, preservando a "naturalidade" de alta fidelidade da destilação avançada.

Fase 1: Moment-Matching Distillation (MMD)

O processo começa destilando um modelo professor base usando Moment-Matching Distillation (MMD). Este método iguala as distribuições intermediárias de denoising ao longo da trajetória de amostragem, gerando um modelo estudante de múltiplos passos que rastreia de perto as marginais do professor. Este estudante é então congelado para servir como uma referência de distribuição estável.

Fase 2: On-Policy Reward Fine-tuning

Na segunda fase, o modelo estudante é ajustado para maximizar uma função de recompensa enquanto mantém a fidelidade distributiva. As principais inovações incluem:

1. Amostragem On-Policy: Em vez de usar pontos de dados off-policy, o método amostra estados ruidosos $x^{pol}_t$ da política atual do estudante (executando o estudante por $K_{student}$ passos, re-ruindo a saída e interrompendo o gradiente). Isso garante que o objetivo de recompensa aproxime diretamente a recompensa esperada do estudante e evita que o desvio distributivo invalide o sinal de treinamento.
2. Regularização Reciclada: A contribuição central é reutilizar a perda de Moment-Matching como um regularizador durante o ajuste fino. Em vez de uma penalidade heurística, esta perda penaliza explicitamente a discrepância entre as distribuições de denoising intermediárias do estudante e as do modelo de referência congelado. Isso fornece um controle interpretável (o peso de regularização) para trocar a maximização da recompensa pela fidelidade distributiva.
3. Gradiente de Passo Único: O método captura toda a faixa de frequência sem o unrolling de múltiplos passos, realizando um único passo de denoising em amostras on-policy corrompidas e avaliando a recompensa.
4. Regularização L2: Um termo L2 adicional penaliza desvios significativos das previsões do modelo inicial destilado via MMD para estabilizar ainda mais o treinamento.

O objetivo total combina o gradiente de recompensa on-policy com o regularizador de moment-matching e o termo L2:
$$\nabla_\theta \mathcal{L}_{RMMD} = \mathbb{E} \left[ \left(\frac{\partial \hat{x}_0}{\partial \theta}\right)^\top \left( m_{student} - m_{teacher} - \lambda \nabla_{\hat{x}_0} R(\hat{x}_0) \right) \right] + \lambda_{reg} \mathcal{L}_{L2}$$

Principais Contribuições

1. Procedimento RMMD: Um novo método para destilação de difusão que simultaneamente maximiza uma função de recompensa, utilizando uma versão on-policy da perda de moment-matching para regularização.
2. Superioridade Empírica: Demonstrou que o RMMD supera consistentemente os baselines de passo único (DI++) e competidores de múltiplos passos (DRaFT, HyperNoise) nas frentes de Pareto FID-Recompensa em diversas funções de recompensa, variando de métricas de nível de pixel a alinhamento semântico (CLIP).
3. Aplicação Científica: Aplicou com sucesso o RMMD ao GenCast, um modelo de previsão meteorológica baseado em difusão de última geração, otimizando o Continuous Ranked Probability Score (CRPS).

Resultados

Experimentos ImageNet

• Vantagem de Multi-Passos: Usar um estudante destilado via MMD de 8 passos como ponto de partida superou significativamente os baselines de 1 passo (DI++). O RMMD de 8 passos alcançou trocas superiores nas frentes de Pareto FID-Recompensa.
• Comparação: O RMMD gerou melhores frentes de Pareto que o DRaFT e HyperNoise em várias resoluções (64x64 e 512x512) e tipos de recompensa (alinhamento CLIP, Inception Score, suavidade).
• Análise Qualitativa: Visualizações mostraram que, enquanto o DRaFT-1 preservou o conteúdo mas introduziu artefatos adversariais, e o DRaFT-2 sofreu com "reward hacking" (desvios drásticos de distribuição), o RMMD integrou modificações sutis para aumentar as recompensas sem desviar significativamente da distribuição original dos dados.

Previsão Meteorológica GenCast

• Aceleração: O modelo destilado alcançou uma aceleração de 7,5× (8 passos vs. 59 avaliações de função) em comparação ao professor.
• Precisão: O modelo RMMD otimizado superou o modelo professor em 93% das variáveis meteorológicas alvo (comparado a 75% para o melhor baseline apenas com MMD).
• Calibração: O RMMD abordou o problema de subdispersão comum na destilação MMD. Embora a variante on-policy tenha mostrado métricas de dispersão ligeiramente piores que a variante off-policy, ela proporcionou as melhores melhorias gerais de CRPS.
• Generalização: A otimização do CRPS para um tempo de antecedência de 12 horas resultou em melhorias consistentes para previsões de até 7 dias, sugerindo que o método melhora a modelagem de distribuição subjacente em vez de apenas sofrer overfitting ao imediato da recompensa.

Significância e Alegações

O artigo afirma que o RMMD fornece uma solução principista para o problema de "drift" inerente à destilação e ao ajuste fino conjuntos. Ao reciclar a perda de moment-matching como um regularizador e utilizar amostragem on-policy, o método permite a otimização de objetivos diferenciáveis complexos (como o CRPS na previsão meteorológica) sem sacrificar a qualidade generativa ou a estabilidade do modelo destilado.

Os autores enfatizam que esta abordagem escala para domínios científicos complexos e de alta dimensão. A aplicação bem-sucedida ao GenCast demonstra que os modelos de difusão podem ser significativamente acelerados e melhorados em suas previsões probabilísticas, abrindo caminho para modelos baseados em difusão mais eficientes e precisos em computação científica onde existem objetivos diferenciáveis. O trabalho sugere que manter uma referência de distribuição estável via moment matching é crucial para um ajuste fino de recompensa eficaz em modelos de difusão.