Resumo Técnico: Modelos de Difusão Ajustados por Metropolis (MADM)

Declaração do Problema

A amostragem a partir de modelos de difusão baseados em pontuação (score) incorre inerentemente em viés decorrente de duas fontes: o erro de aproximação da função de pontuação aprendida e o erro de discretização temporal introduzido por integradores numéricos. No framework padrão Predictor-Corrector, uma etapa de preditor (frequentemente uma EDO de Fluxo de Probabilidade) move uma amostra entre níveis de ruído, seguida por uma etapa de corretor (tipicamente o Algoritmo de Langevin Não Ajustado, ULA) para refinar a amostra no nível de ruído atual.

Embora o ULA utilize a função de pontuação para propor movimentos, ele é, em si mesmo, um amostrador enviesado porque discretiza um processo de difusão de Langevin contínuo. Consequentemente, o "corretor" no esquema Predictor-Corrector introduz seu próprio viés de discretização, falhando em deixar a distribuição marginal alvo invariante. Mecanismos padrão de aceitação-rejeição Metropolis-Hastings (MH) ou Barker, que corrigem tal viés em Cadeias de Markov Monte Carlo (MCMC), não podem ser aplicados diretamente a modelos de difusão. Isso ocorre porque o MH exige a avaliação da razão de densidades alvo, $p_t(\tilde{x})/p_t(x)$, o que é intratável em modelos de difusão onde apenas a função de pontuação $s_\phi(x, t) \approx \nabla_x \log p_t(x)$ está disponível.

Metodologia

Os autores propõem Modelos de Difusão Ajustados por Metropolis (MADM), um framework que substitui corretores ULA enviesados por etapas de Langevin corrigidas por aceitação-rejeição. A inovação central é um método para calcular a probabilidade de aceitação necessária usando apenas a função de pontuação, contornando a necessidade da razão de densidades intratável.

1. Identidade da Razão de Densidades

Os autores estabelecem que o logaritmo da razão de densidades alvo pode ser expresso como uma integral de linha da função de pontuação ao longo do caminho entre o estado atual $x$ e a proposta $\tilde{x}$:
$$\log \left( \frac{p_t(\tilde{x})}{p_t(x)} \right) = \int_0^1 \langle s(x + u(\tilde{x} - x), t), \tilde{x} - x \rangle \, du$$
Esta identidade permite a estimativa da razão de densidades usando apenas avaliações da função de pontuação.

2. Correção Exata via Fábrica de Bernoulli

Para alcançar uma correção exata sem avaliar a densidade, os autores propõem um algoritmo baseado em uma fábrica de Bernoulli de duas moedas (Gonçalves et al., 2017).

• Pressuposto: Existe um limite $C(x, \tilde{x})$ tal que o valor absoluto do integrando na integral de linha é limitado por $C$.
• Mecanismo: O algoritmo constrói um estimador não enviesado $W$ do exponencial da integral (escalado pelo limite) usando uma série de potências truncada por Poisson.
• Procedimento: O algoritmo utiliza um procedimento de "duas moedas" para simular uma decisão de aceitação-rejeição de Barker. Primeiro, rejeita com uma probabilidade fixa $\alpha'$, e se não for rejeitado, lança uma moeda com probabilidade $W$. Se a moeda for favorável, a proposta é aceita; caso contrário, a rodada recomeça.
• Resultado: Este procedimento produz um corretor de Langevin ajustado por Barker exato que deixa a distribuição alvo $p_t$ invariante, desde que a pontuação seja exata e o pressuposto do limite seja válido.

3. Aproximação Eficiente via Regra de Simpson

Reconhecendo que o método exato de fábrica de Bernoulli requer um número aleatório e potencialmente grande de avaliações de pontuação, os autores propõem uma aproximação escalável.

• Método: A integral de linha é aproximada deterministicamente usando a Regra 1/3 de Simpson (uma regra de quadratura de Newton-Cotes). Isso requer avaliar a pontuação no estado atual, na proposta e no ponto médio.
• Aceitação: A estimativa resultante é inserida na probabilidade de aceitação padrão de Metropolis-Hastings.
• Propriedades: Embora isso introduza um pequeno viés devido à truncagem da quadratura, a aproximação alcança uma precisão de ordem $O(h^{5/2})$ no tamanho do passo $h$. É computacionalmente barata, exigindo apenas uma avaliação adicional de pontuação no ponto médio por proposta.

Contribuições Principais

1. Framework MADM: Introdução de um framework de corretor ajustado por Metropolis para modelos de difusão que substitui sistematicamente etapas ULA enviesadas.
2. Método Exato: O primeiro método exato para ajustar correções de Langevin em modelos de difusão usando uma fábrica de Bernoulli de duas moedas, exigindo apenas a função de pontuação e um limite na integral de linha.
3. Aproximação Escalável: Uma aproximação prática via Regra de Simpson que alcança alta precisão de ordem ($O(h^{5/2})$) com custo marginal próximo de zero (uma avaliação extra de pontuação).
4. Validação Empírica: Demonstração de melhoria na qualidade das amostras em conjuntos de dados sintéticos e de imagem.

Resultados

Os autores avaliam o MADM em conjuntos de dados sintéticos 2D e benchmarks padrão de imagem (CIFAR-10, FFHQ, AFHQv2, ImageNet-64) usando modelos pré-treinados de Karras et al. (2022).

• Conjuntos de Dados Sintéticos: Em conjuntos de dados 2D (espiral, funil, triângulo de Sierpinski, cata-vento), o amostrador ancestral padrão (equivalente a um preditor com um corretor ULA) produziu amostras com outliers significativos fora do suporte da distribuição de dados verdadeira. As etapas corretoras do MADM moveram com sucesso essas amostras de volta para o suporte dos dados, eliminando outliers.
• Conjuntos de Dados de Imagem:
  • Para modelos incondicionais (CIFAR-10, FFHQ, AFHQv2), aplicar uma etapa corretora MADM por etapa de preditor produziu melhorias consistentes, embora modestas, na Distância de Inception Fréchet (FID) em comparação com bases determinísticas (solucionadores de EDO Heun e Euler).
  • Para o modelo class-conditional ImageNet-64, os ganhos foram mais pronunciados. Usar MADM com um preditor Euler produziu maiores melhorias no FID do que usar um preditor Heun de ordem superior sem correção. Isso sugere que os ajustes MH são mais eficazes em mitigar erros de discretização do que simplesmente aumentar a ordem do solucionador de EDO quando o número de passos é grande.
• Custo Computacional: O método exato (Algoritmo 2) mostrou-se computacionalmente caro devido ao número aleatório de avaliações de pontuação requeridas. Consequentemente, os experimentos utilizaram principalmente a aproximação pela Regra de Simpson (Algoritmo 3), que forneceu ganhos consistentes com sobrecarga gerenciável.

Significado e Alegações

O artigo alega fornecer o primeiro método exato para ajustar correções de Langevin em modelos de difusão, abordando o viés inerente dos corretores ULA. Ao aproveitar a função de pontuação para estimar razões de densidades via integrais de linha, o MADM permite o uso de critérios rigorosos de aceitação MCMC em modelagem generativa sem exigir acesso à densidade de dados intratável.

Os autores posicionam o MADM como uma extensão natural do framework Predictor-Corrector que melhora a qualidade das amostras. Eles observam que, embora o método exato seja teoricamente significativo, a aproximação pela Regra de Simpson oferece uma solução prática e escalável que melhora consistentemente o FID em benchmarks de imagem. O trabalho assume uma estimativa de pontuação perfeita, reconhecendo isso como uma limitação inerente dos modelos de difusão e não como uma falha específica de seu método. Sugere-se trabalho futuro focado no desenvolvimento de variantes eficientes do método exato ou de suas sub-rotinas.