Resumo Técnico: URGE – Aproximação Simples e Inferência Livre de Derivadas para Escalonamento em Tempo de Inferência de Modelos de Difusão via Monte Carlo Sequencial em Medidas de Trajetória

1. Formulação do Problema

Modelos de difusão modernos e métodos generativos baseados em fluxo são tipicamente treinados para capturar a distribuição de dados $p_{\text{data}}$. No entanto, a implantação frequentemente exige direcionar a geração para objetivos específicos de downstream (por exemplo, validade física, alinhamento texto-imagem) sem retreinamento. Isso é formalizado como amostragem de uma posteriori inclinada por recompensa $q(x) \propto p_{\text{data}}(x)r(x)$.

Os métodos atuais de escalonamento em tempo de inferência, particularmente o guia (guidance), tentam alcançar isso modificando a deriva da EDE (Equação Diferencial Estocástica) de amostragem com um termo de gradiente derivado de um potencial de guia $G$. Embora empiricamente eficazes, o guia introduz uma limitação matemática fundamental: não garante amostragem da distribuição inclinada exata $q(x)$. A amostragem exata requer a transformação $h$ de Doob, que envolve resolver uma equação de Kolmogorov reversa e é frequentemente intratável em altas dimensões.

Alternativas existentes, como FK-Corrector e AFDPS, tentam corrigir o viés da geração guiada usando Monte Carlo Sequencial (SMC) por meio de reponderação de partículas. No entanto, esses métodos dependem de geradores infinitesimais que exigem acesso a derivadas de ordem superior da função de recompensa (especificamente o escore $\nabla \log p_t$, o Laplaciano $\Delta r$ e o Hessiano do guia). Esses requisitos acarretam uma sobrecarga computacional significativa e tornam esses métodos inaplicáveis a modelos de recompensa de caixa preta (por exemplo, modelos de recompensa neural usados em geração de texto para imagem) onde as derivadas estão indisponíveis ou são caras.

O desafio central abordado por este trabalho é: Podemos projetar um procedimento de escalonamento em tempo de inferência livre de aproximações que elimine o viés do guia subótimo sem exigir informações de derivada da recompensa?

2. Metodologia: URGE

Os autores propõem o URGE (Reamostragem Não Viciada via Estimativa de Girsanov), um algoritmo de escalonamento em tempo de inferência livre de derivadas que realiza reamostragem sobre trajetórias de geração (espaço de trajetórias) em vez de amostras individuais geradas (espaço de partículas).

Mecanismo Central

O URGE aproveita o teorema de Girsanov para realizar uma mudança de medida no espaço de trajetórias do processo de difusão.

1. Geração Guiada: Múltiplas trajetórias $\{X^{(i),G}_t\}$ são simuladas usando uma EDE guiada com um potencial de guia arbitrário $G$ (que pode ser imperfeito).
   $$dX^G_t = (v(X^G_t, t) + V(t)^2 \nabla_x G(X^G_t, t))dt + V(t)dW_t$$
2. Reponderação por Trajetória: Em vez de calcular pesos baseados em geradores infinitesimais (que exigem derivadas), o URGE atribui um peso multiplicativo a cada trajetória inteira com base na derivada de Radon-Nikodym entre a medida de trajetória inclinada por recompensa alvo $Q$ e a medida de trajetória guiada $P_G$.
   O peso de importância para uma trajetória sobre $[s, t]$ é:
   $$w^{\text{URGE}}_{s,t} = \exp\left( \int_s^t -V(\tau)\nabla_x G(X^G_\tau, \tau)^\top dW_\tau - \frac{1}{2}\int_s^t V(\tau)^2 \|\nabla_x G(X^G_\tau, \tau)\|^2 d\tau + \int_s^t dr(X^G_\tau, \tau) \right)$$
   Crucialmente, esta formulação não requer $\nabla r$, $\Delta r$ ou $\nabla \log p_t$. Requer apenas a função de recompensa $r$ e o guia $G$.
3. Reamostragem: Em passos de tempo discretos, as partículas são reamostradas proporcionalmente aos seus pesos de trajetória. Isso implementa um mecanismo de "nascimento-morte" onde trajetórias de alta recompensa são replicadas e as de baixa recompensa são eliminadas, impulsionando a distribuição empírica em direção à posteriori verdadeira $q$.

Equivalência Teórica

O artigo estabelece uma equivalência teórica entre a reamostragem no espaço de trajetórias (URGE) e a reamostragem no espaço de partículas (FK-Corrector/AFDPS).

• Os autores provam que o gerador marginalizado do processo de espaço de trajetórias do URGE, quando condicionado ao estado terminal, recupera o termo exato de reponderação infinitesimal usado pelo FK-Corrector e AFDPS.
• Isso confirma que o URGE é livre de aproximações: ele converge para a mesma distribuição alvo que os métodos no espaço de partículas teoricamente ótimos (mas pesados em derivadas), mas alcança isso elevando a construção de pesos ao espaço de trajetórias, contornando assim a necessidade de cálculos explícitos de derivadas.

3. Contribuições Principais

1. Novo Framework de Mistura de Especialistas: Os autores introduzem um método que integra SMC com processos generativos por meio de reamostragem sobre trajetórias em vez de amostras. Isso permite uma construção de pesos flexível via mudança de medida de Girsanov sem exigir derivadas de recompensa.
2. Implementação Livre de Derivadas: O URGE elimina a necessidade de funções de escore, Hessianos ou Laplacianos da recompensa. Isso o torna aplicável a modelos de recompensa de caixa preta (por exemplo, CLIP, ImageReward, modelos ajustados por DPO) onde os gradientes são inacessíveis.
3. Equivalência Teórica: O artigo prova que a reponderação por trajetória do URGE, quando marginalizada, recupera exatamente a estrutura de ponderação dos métodos baseados em derivadas de última geração (FK-Corrector/AFDPS), garantindo que ambos os esquemas produzam a mesma lei terminal livre de aproximações.
4. Superioridade Empírica: O URGE supera consistentemente as linhas de base existentes (incluindo variantes de FK-Steering e AFDPS) em tarefas sintéticas e do mundo real, sendo significativamente mais simples de implementar e totalmente livre de gradientes.

4. Resultados Experimentais

Os autores avaliam o URGE em três benchmarks:

• Modelo de Mistura Gaussiana (Sintético):

  • O URGE alcança o menor erro em todas as métricas (MMD, SWD, Média L2, Frob de Covariância) em comparação com Guia Puro, AFDPS e FK-Steering.
  • Demonstra eficiência de convergência superior, alcançando melhores resultados com menos partículas do que as linhas de base.
  • O URGE permanece robusto em dimensões variáveis e passos de discretização, mostrando menor sensibilidade ao tamanho do passo do que os métodos concorrentes.

• Problemas Inversos (Restauração de Imagem):

  • Testado em Desembaçamento Gaussiano, Desembaçamento por Movimento, Super Resolução e Preenchimento de Caixas usando ImageNet-256 e FFHQ-256.
  • O URGE alcança desempenho comparável ao AFDPS (que requer derivadas) e supera significativamente outras linhas de base em PSNR e LPIPS.
  • O URGE demonstra melhor escalabilidade computacional, alcançando tempos de inferência mais curtos que o AFDPS-ODE enquanto mantém alta qualidade de reconstrução.

• Geração de Texto para Imagem:

  • Testado no Stable Diffusion v1.5 e SDXL com variantes ajustadas por DPO.
  • O URGE supera o FK-Steering e o modelo base em qualidade estética (ImageReward, HPS) e fidelidade ao prompt (GenEval, CLIP-Score).
  • Notavelmente, o URGE permite que um modelo menor (SD v1.5) iguale ou se aproxime do desempenho de um modelo maior (SDXL) ao usar escalonamento de partículas, enquanto o FK-Steering mostra retornos decrescentes.
  • O URGE exibe melhoria monotônica da recompensa com o aumento da contagem de partículas, ao contrário do FK-Steering.

5. Significado e Alegações

O artigo alega que o URGE fornece um framework de amostragem de posteriori baseado em ensemble, livre de treinamento e livre de aproximações. Seu significado principal reside em unificar estratégias de ponderação baseadas em trajetórias e baseadas em estados. Ao elevar o mecanismo de reamostragem do espaço de partículas para o espaço de trajetórias, o URGE:

• Remove o gargalo das derivadas: Permite o uso de modelos de recompensa complexos e de caixa preta na orientação de difusão sem o custo computacional de calcular derivadas de alta ordem.
• Garante correção teórica: Garante convergência para a posteriori verdadeira inclinada por recompensa, corrigindo o viés inerente dos métodos de guia padrão.
• Oferece eficiência prática: Alcança qualidade de geração de última geração com implementação mais simples e sobrecarga computacional competitiva, tornando-o adequado para tarefas generativas em grande escala e do mundo real.

Os autores enfatizam que, embora trabalhos anteriores (FK-Corrector, AFDPS) fornecessem insights teóricos sobre a correção do guia, eles eram limitados pela sua dependência de informações de derivada. O URGE preenche essa lacuna, oferecendo uma solução prática e livre de derivadas que retém as garantias teóricas da amostragem livre de aproximações.