Resumo Técnico: Aprendendo a Amostrar de Modelos de Difusão via Aprendizado por Reforço Inverso

Declaração do Problema

Os modelos de difusão tornaram-se um padrão para a geração de dados de alta dimensão, mas o procedimento de amostragem continua sendo um gargalo crítico. Embora a rede neural de denoising seja tipicamente fixa após o treinamento, o algoritmo de amostragem em si requer um ajuste cuidadoso de hiperparâmetros (ex: cronogramas de ruído, escalas de orientação, perfis de estocasticidade). Atualmente, esse ajuste é realizado via buscas de grade empíricas custosas. Para modelos complexos como aqueles treinados no ImageNet-64, o custo computacional de uma busca exaustiva de hiperparâmetros pode se aproximar ou até exceder o custo de treinamento do próprio denoiser. Além disso, abordagens de aprendizado por reforço (RL) ou de controle existentes para modelos de difusão frequentemente exigem funções de recompensa explícitas ou o ajuste fino do denoiser, o que contradiz o objetivo de otimizar o amostrador sem modificar o gerador pré-treinado.

Metodologia

Os autores propõem enquadrar a amostragem de difusão como um problema de tomada de decisão sequencial solucionável via Aprendizado por Reforço Inverso (IRL). A inovação central é aprender uma política de amostragem que corresponda às distribuições de estados alvo sem acesso a trajetórias de especialistas, recompensas explícitas ou retreinamento do denoiser.

1. Formulação MDP

O processo de amostragem de difusão é formalizado como um Processo de Decisão de Markov (MDP) de horizonte finito:

• Estado ($s_t$): Definido como $(x_t, \sigma_t)$, onde $x_t$ é a amostra atual e $\sigma_t$ é o nível de ruído.
• Ação ($a_t$): Decisões de controle discretas aplicadas à dinâmica de amostragem, como selecionar escalas de orientação, níveis de injeção de ruído ou decidir se deve realizar o renoising (reiniciar) para um nível de ruído anterior.
• Dinâmica: O kernel de transição é induzido pelas dinâmicas de difusão discretizadas e pela ação de controle selecionada. O processo termina quando o nível de ruído atinge $\Sigma_0$ (um estado absorvente).

2. Aprendizado por Reforço Inverso Baseado Apenas em Estados

Ao contrário do IRL tradicional, que frequentemente assume ações ou trajetórias de especialistas, este método opera sob um regime de apenas estados. O "especialista" é definido unicamente pelas distribuições marginais de amostras de dados em níveis de ruído selecionados.

• Objetivo: O objetivo é encontrar uma política $\pi_\theta$ cuja medida de ocupação de estado $\mu_\theta$ (a distribuição marginal dos estados visitados) corresponda à medida de ocupação do especialista $\mu_E$.
• Otimização: Em vez de aprender uma função de recompensa, os autores minimizam diretamente a divergência-$f$ entre as medidas de ocupação do especialista e da política:
  $$\min_\theta D_f(\mu_E \parallel \mu_\theta)$$
  O artigo explora especificamente as divergências de Kullback-Leibler (KL) e Reverse-KL (rKL), observando que a escolha da divergência-$f$ dita o compromisso entre cobertura de modo (diversidade) e busca de modo (fidelidade).

3. Gradiente de Política e Otimização

Os autores derivam um teorema de gradiente de política para minimizar divergências-$f$ diretamente. O gradiente é expresso como:
$$\nabla_\theta L_f(\theta) = \mathbb{E}_{s_{1:T} \sim p_\theta} \left[ \frac{1}{T} \sum_{t=1}^T \nabla_\theta \log \pi_\theta(a_t|s_t) \sum_{t' \ge t} h_f\left(\frac{\mu_E(s_{t'})}{\mu_\theta(s_{t'})}\right) \right]$$
onde $h_f$ é uma função derivada do gerador da divergência-$f$.

• Estimativa de Razão de Densidade: Como $\mu_E$ e $\mu_\theta$ são desconhecidos, a razão de densidade é estimada usando um discriminador treinado para distinguir entre estados do especialista (dados com ruído) e estados gerados pela política.
• Estabilização: Para lidar com a alta variância nas estimativas de gradiente, o método emprega um objetivo clipado estilo PPO, correções de amostragem de importância e uma linha de base de média de lote global.

4. Estratégias de Amostragem Avaliadas

O framework é aplicado a três intervenções específicas de tempo de inferência:

1. Orientação de Classe Livre Adaptativa (Adaptive Classifier-Free Guidance): Aprendendo uma escala de orientação dependente do estado $\omega(x, \sigma)$.
2. Injeção de Estocasticidade: Aprendendo um nível de injeção de ruído dependente do estado $\gamma_{EDM}(x, \sigma)$.
3. Renoising (Amostragem de Reinício): Aprendendo quando reverter o processo para um nível de ruído anterior para escapar de gerações locais ruins.

Principais Contribuições

1. Formulação MDP para Amostragem: O artigo enquadra a amostragem de difusão como um MDP para adaptar amostradores no tempo de inferência sem retreinar o denoiser.
2. Correspondência Direta de Marginais de Estado: Deriva um objetivo de gradiente de política que minimiza divergências-$f$ entre medidas de ocupação de estado, contornando a necessidade de modelagem de recompensa explícita ou dados de ação de especialistas.
3. Compromissos Interpretáveis: O framework expõe um compromisso controlável entre diversidade de amostra (Recall) e fidelidade (Precision) através da escolha da divergência-$f$ (KL vs. rKL).
4. Controle Dependente do Estado: As políticas aprendidas atuam como controles de tempo de inferência que intervêm seletivamente (ex: aplicando alta orientação apenas em níveis de ruído específicos ou para tipos de amostras específicos) em vez de aplicar cronogramas globais fixos.

Resultados Experimentais

O método foi avaliado no CIFAR-10, FFHQ e ImageNet-64 usando modelos EDM pré-treinados.

• Desempenho: As políticas aprendidas igualam ou superam o desempenho das melhores linhas de base ajustadas manualmente (encontradas via busca de grade) em termos de métricas FID, Precisão e Recall.
  • Orientação Adaptativa: Melhorou o compromisso qualidade-diversidade em comparação com a orientação fixa. A divergência KL melhorou o Recall (diversidade), enquanto a rKL melhorou a Precisão (fidelidade).
  • Injeção de Estocasticidade: Os perfis dependentes do estado aprendidos recuperaram os padrões qualitativos dos cronogramas ótimos do EDM (baixa estocasticidade ao final da trajetória) e superaram as linhas de base fixas em conjuntos de dados complexos como o ImageNet.
  • Renoising: Reduziu significativamente o Número de Avaliações de Função (NFE) no ImageNet, mantendo ou até melhorando o FID em comparação com estratégias de reinício padrão.
• Eficiência Computacional:
  • Treinamento: Uma única execução de treinamento substitui a busca de grade exaustiva. No ImageNet-64, isso reduz o orçamento de otimização em até 9$\times$ (de ~1437 EFLOPs para uma busca de grade de 3 hiperparâmetros para ~161 EFLOPs para o treinamento de IRL).
  • Inferência: A política adiciona um overhead modesto de 16% a 40% em FLOPs em relação ao denoiser sozinho, o que é considerado insignificante dado o fim do ajuste manual.
• Descobertas de Ablação:
  • Condicionamento: Políticas condicionadas ao estado completo $(x, \sigma)$ superam significativamente aquelas condicionadas apenas ao nível de ruído $\sigma$, e performam quase tão bem quanto aquelas com rótulos de classe explícitos $(x, \sigma, c)$, sugerendo que a política aprende implicitamente características discriminativas de classe a partir do conteúdo da imagem.
  • Escolha de Divergência: KL e rKL mostraram-se as mais eficazes para controlar o compromisso precisão-recall, enquanto divergências simétricas (TV, Hellinger) produziram resultados menos interpretáveis.

Significância e Alegações

O artigo afirma que este framework fornece uma alternativa sistemática e automatizada ao ajuste manual de hiperparâmetros para a amostragem de difusão. Ao tratar a amostragem como um problema de IRL, ele permite a otimização de estratégias de amostragem complexas e dependentes do estado que são difíceis de projetar manualmente.

Os autores enfatizam que o método é particularmente valioso para modelos de grande escala onde o custo de avaliação por tentativa em uma busca de grade é proibitivo. A abordagem não modifica o modelo generativo em si, mas aprende uma "política de controle" ótima para a trajetória de inferência, ofereando um caminho prático para melhorar a qualidade e a diversidade das amostras sem o fardo computacional de retreinamento ou busca exaustiva. O artigo permanece modesto, observando que os benefícios dependem do conjunto de dados e que o framework atual depende de espaços de ação discretos, sugerindo trabalhos futuros em espaços de ação contínuos e inicialização mais robusta.