Likelihood Matching for Diffusion Models

Resumo Técnico: Correspondência de Verossimilhança para Modelos de Difusão

Enunciado do Problema
Os modelos de difusão atuais predominantes, particularmente os Modelos Generativos Baseados em Score (SGMs), baseiam-se na Correspondência de Score (SM - Score Matching) como seu principal objetivo de treinamento. Embora empiricamente eficaz, a SM minimiza a discrepância entre uma função de score parametrizada e a verdadeira função de score em vários níveis de ruído. Teoricamente, esta abordagem minimiza apenas um limite superior da log-verossimilhança negativa da distribuição de dados original, em vez da própria verossimilhança. Essa conexão indireta pode levar a uma perda de eficiência estatística em comparação com a Estimativa de Máxima Verossimilhança (MLE), que é conhecida por produzir estimadores consistentes e eficientes. Além disso, a SM padrão foca exclusivamente na informação de momento de primeira ordem (o score/gradiente), potencialmente negligenciando informações estruturais de ordem superior contidas nas densidades de transição.

Metodologia
Os autores propõem a Correspondência de Verossimilhança (LM - Likelihood Matching), um framework que visa diretamente a verossimilhança da distribuição de dados ao aproveitar a equivalência entre a verossimilhança dos dados alvo e a verossimilhança do caminho do processo de difusão reversa.

1. Fundamentação Teórica:
   O método baseia-se na propriedade de que a verossimilhança conjunta do caminho de difusão reversa é equivalente à verossimilhança da distribuição de dados original (até constantes relacionadas ao processo direto). Os autores estabelecem que maximizar a verossimilhança exata do caminho da difusão reversa é equivalente a maximizar a log-verossimilhança dos dados.

2. Estimativa de Quasi-Máxima Verossimilhança (QMLE):
   Como as densidades de transição reversas exatas são intratáveis, os autores as aproximam usando uma abordagem de Quasi-Máxima Verossimilhança. Eles modelam cada densidade de transição reversa como uma distribuição Gaussiana onde a média condicional e a covariância são ajustadas à verdadeira função reversa.

   • Insight Chave: A média condicional e a covariância do processo reverso dependem não apenas da função de score ($\nabla \log q_t$), mas também da função Hessiana ($\nabla^2 \log q_t$).
   • Parametrização: Os autores parametrizam tanto o score quanto a Hessiana usando redes neurais (especificamente U-Nets). Para gerenciar a complexidade computacional, a Hessiana é frequentemente aproximada usando uma estrutura de diagonal mais baixo posto (diagonal-plus-low-rank).

3. O Objetivo de Correspondência de Verossimilhança:
   O objetivo proposto, a Correspondência de Verossimilhança (LM), é derivado pela minimização da quasi-log-verossimilhança negativa das transições reversas.

   • O objetivo unifica a correspondência de score e a correspondência de covariância.
   • Ele incorpora um mecanismo de ponderação de verossimilhança implícita via matriz de covariância ajustada $(I + \sigma^2 H_t)^{-1/2}$, que ajusta automaticamente o peso do termo de correspondência de score com base na curvatura local.
   • O treinamento utiliza um amostrador estocástico que seleciona aleatoriamente passos de tempo ao longo da trajetória de difusão para computar a perda, garantindo o uso eficiente da informação evolutiva do SDE direto.

4. Amostrador Estocástico:
   Um novo procedimento de amostragem é introduzido utilizando tanto a informação de score quanto de Hessiana aprendida. O passo reverso é definido como:
   $$\tilde{Y}_{t-1} = \hat{\mu}_{t-1|t}(\tilde{Y}_t) + \hat{\Sigma}_{t-1|t}^{1/2}(\tilde{Y}_t) Z_t$$
   onde a média $\hat{\mu}$ e a covariância $\hat{\Sigma}$ são funções do score e da Hessiana estimados, respectivamente.

Principais Contribuições

• Novo Objetivo de Treinamento: O artigo introduz um objetivo de treinamento de quasi-verossimilhança (LM) que aproxima a log-verossimilhança do caminho reverso, combinando efetivamente a correspondência de score e a correspondência de covariância com ponderação de verossimilhança implícita.
• Amostragem Aprimorada pela Hessiana: Um amostrador estocástico é derivado aproveitando tanto a informação de score quanto de Hessiana, oferecendo uma aproximação mais precisa da densidade de transição reversa do que os amostradores padrão do tipo DDPM.
• Garantias Teóricas:
  • Convergência: Garantias de convergência não assintótica são fornecidas em distância de variação total. O limite de erro escala como $O(d^3 \log^{4.5} T / T)$ para o erro do passo reverso, com termos adicionais para erros de estimativa de score e Hessiana.
  • Consistência: O artigo prova a consistência da proposta QMLE sob um cenário paramétrico órfão (oracle), mostrando que o estimador converge para os verdadeiros parâmetros dos dados conforme o tamanho da amostra e os passos de difusão aumentam.
• Implementação Eficiente: O uso de aproximações de Hessiana de baixo posto e a fórmula de Sherman-Morrison-Woodbury permitem o cálculo eficiente da verossimilhança e dos passos de amostragem.

Resultados Experimentais
Os autores validam o método em conjuntos de dados sintéticos e do mundo real:

• Dados Sintéticos: Em misturas de distribuições Gaussianas e t, o LM superou consistentemente a Correspondência de Score (SM) padrão, particularmente ao capturar modos bem separados. Em uma distribuição Swiss-roll altamente enrolada (representando variedades curvas), o LM demonstrou desempenho superior, validando o benefício da covariância informada pela Hessiana.
• Estimativa de Parâmetros: Em um cenário paramétrico de mistura Gaussiana, o LM alcançou menor Erro Absoluto Médio (MAE) e erro padrão na estimativa de parâmetros em comparação com a SM, apoiando as alegações de consistência teórica.
• Conjuntos de Dados de Imagens (MNIST, CIFAR-10, CelebA):
  • O LM alcançou menor Distância de Fréchet Inception (FID) e melhores pontuações de Log-Verossimilhança Negativa (NLL) em comparação com as bases de SM.
  • O desempenho melhorou com o posto da aproximação da Hessiana ($r$), atingindo o pico em torno de $r=20-30$.
  • Compromisso Velocidade-Qualidade: O LM alcançou um FID menor para um tempo de execução (wall-clock time) comparável e exigiu menos passos de amostragem reversa para atingir alta qualidade visual. Por exemplo, no MNIST, o LM produziu dígitos reconhecíveis após apenas 20 iterações reversas, enquanto as amostras da SM permaneceram borradas.

Significância e Alegações
O artigo afirma que a Correspondência de Verossimilhança oferece uma alternativa estatisticamente mais rigorosa à Correspondência de Score ao visar diretamente a verossimilhança dos dados através de um framework de quasi-verossimilhança. Ao incorporar informação de segunda ordem (a Hessiana), o método captura densidades de transição mais ricas e melhora a eficiência da amostragem, particularmente em regimes de poucos passos e em variedades de dados complexas e curvas. Os autores enfatizam que, embora o LM incorra em custos computacionais mais altos por passo devido à modelagem da Hessiana, a redução no número de passos de amostragem necessários e a melhoria na qualidade da geração (FID e NLL) validam sua eficácia para o treinamento de modelos generativos de alta fidelidade. O trabalho preenche a lacuna entre a eficiência estatística da MLE e o sucesso prático dos modelos de difusão.