Proximal Diffusion Neural Sampler

Resumo Técnico: Amostrador Neural de Difusão Proximal (PDNS)

Declaração do Problema
Aprendizado de um amostrador neural baseado em difusão para gerar amostras de uma distribuição alvo não normalizada (por exemplo, uma distribuição de Boltzmann $\pi(x) \propto e^{-\beta E(x)}$) é uma tarefa fundamental em estatística computacional, inferência bayesiana e física estatística. Este problema é frequentemente formulado como um problema de Controle Estocástico Ótimo (SOC) em medidas de trajetória, onde uma rede neural parametriza um processo de controle para transportar uma distribuição de referência simples até o alvo.

No entanto, o treinamento desses amostradores enfrenta desafios significativos quando a distribuição alvo é multimodal com barreiras de energia elevadas separando os modos. Em tais cenários, métodos padrão baseados em SOC (incluindo variantes de Entropia Cruzada e Entropia Cruzada de Desruído Ponderada) frequentemente sofrem de colapso de modos. No início do treinamento, o modelo captura apenas alguns modos; a função de perda global então reforça esses modos capturados com base nas trajetórias limitadas geradas, fazendo com que a massa de probabilidade se concentre em um subconjunto de modos enquanto ignora outros. Isso leva a uma otimização instável e a uma cobertura pobre da distribuição alvo.

Metodologia
Os autores propõem o Amostrador Neural de Difusão Proximal (PDNS), um framework unificado projetado para abordar o colapso de modos reformulando o processo de treinamento como uma sequência de subproblemas locais em vez de uma única minimização global.

1. Framework Unificado: O artigo unifica primeiro amostradores de difusão contínuos e discretos sob um único framework SOC. Ele define a medida de trajetória ótima $P^*$ como a solução para minimizar um objetivo de entropia relativa sujeito a uma medida de trajetória de referência $P_{ref}$ que satisfaz uma condição "sem memória" (independência entre distribuições inicial e terminal).
2. Método do Ponto Proximal: Para mitigar o colapso de modos, o PDNS aplica o método do ponto proximal ao espaço de medidas de trajetória. Em vez de otimizar diretamente o objetivo global, o algoritmo resolve iterativamente um subproblema local no passo $k$:
   $$P_{\theta^*_k} = \arg\min_{P_\theta} \left( -E_{P_\theta}[r(X_T)] + KL(P_\theta \| P_{ref}) + \frac{1}{\eta_k} KL(P_\theta \| P_{\theta_{k-1}}) \right)$$
   Aqui, $\eta_k$ é um tamanho de passo proximal, e o termo $KL(P_\theta \| P_{\theta_{k-1}})$ atua como um regularizador mantendo a iteração atual próxima à anterior.
3. Propriedades Teóricas: Os autores provam que a solução ótima para o $k$-ésimo passo proximal é uma interpolação geométrica entre a iteração anterior e a medida de trajetória ótima. Especificamente, a distribuição terminal do $k$-ésimo passo, $P_k$, é uma interpolação geométrica entre a distribuição terminal de referência $\nu$ e o alvo $\pi$. À medida que $k \to \infty$, a sequência de medidas de trajetória converge para o alvo $P^*$. Este procedimento em etapas garante que o modelo explore o espaço de estados gradualmente, prevenindo o compromisso prematuro com um único modo.
4. Instanciação Prática (Proximal WDCE): Para eficiência, os autores instanciam o PDNS usando um objetivo de Entropia Cruzada de Desruído Ponderada Proximal (Proximal WDCE). Esta abordagem evita armazenar trajetórias completas usando pesos de importância acumulados online.
   • Caso Contínuo: O objetivo envolve corresponder uma pontuação de desruído (ou ponte) com um termo de peso derivado do tamanho de passo proximal $\eta_k$ e da função de energia.
   • Caso Discreto: O objetivo adapta a difusão discreta mascarada, utilizando um esquema de ponderação similar baseado na derivada de Radon-Nikodym das medidas de trajetória.
5. Design do Agendador: O artigo discute estratégias para selecionar o tamanho de passo $\eta_k$ (ou seu equivalente $\lambda_k$). Um $\eta_k$ menor (regularização mais forte) preserva a cobertura de modos, mas desacelera a convergência, enquanto um $\eta_k$ maior arrisca o colapso de modos. Os autores propõem tanto agendamentos predefinidos ($\lambda_k$ decrescente linearmente) quanto uma estratégia adaptativa que ajusta $\eta_k$ com base em uma restrição de região de confiança na divergência KL entre o modelo atual e o próximo alvo.

Principais Contribuições

1. Framework Unificado: O artigo integra amostradores baseados em SOC contínuos e discretos em uma única formulação de medida de trajetória.
2. Amostrador Neural de Difusão Proximal (PDNS): Introduz um framework inovador que aproveita iterações de ponto proximal no espaço de trajetórias para estabilizar o aprendizado e mitigar o colapso de modos em configurações multimodais.
3. Proximal WDCE: Deriva instanciamações concretas e computacionalmente eficientes do PDNS para domínios contínuos e discretos usando perdas de entropia cruzada de desruído ponderada proximal.
4. Estratégias de Agendamento: Os autores delineiam estratégias principiantes para selecionar tamanhos de passo proximais, incluindo um agendador adaptativo que equilibra velocidade de convergência e estabilidade.

Resultados Experimentais
Os autores avaliam o PDNS em uma ampla gama de benchmarks contínuos e discretos:

• Tarefas Contínuas:
  • Funções de Energia Sintéticas: O PDNS supera bases fortes (incluindo SMC, DDS, PIS e AS) em benchmarks multimodais como o potencial de 32 poços (MW54), Funil e misturas gaussianas de alta dimensão (GMM40). Alcança resultados state-of-the-art na mistura t-Student (MoS) e potenciais baseados em partículas (sistemas de Lennard-Jones LJ-13 e LJ-55).
  • Dinâmica Molecular: No sistema de Dipeptídeo de Alanina (AD), o PDNS alcança divergências KL comparáveis ao método state-of-the-art ASBS e captura com precisão estruturas conformacionais em gráficos de Ramachandran.
• Tarefas Discretas:
  • Física Estatística: O PDNS supera significativamente LEAPS e Metropolis-Hastings (MH) na amostragem de modelos de Ising e Potts em baixas temperaturas (perto ou abaixo de pontos críticos), onde a separação de modos é severa. Alcança erros menores em magnetização e correlações de 2 pontos.
  • Otimização Combinatória: Aplicado ao problema do Corte Máximo em grafos Barabási–Albert e Erdős–Rényi, o PDNS demonstra a capacidade de alcançar qualidades de solução comparáveis a solucionadores de verdade (Gurobi).

Significado e Alegações
O artigo alega que o PDNS fornece uma solução robusta para o problema de colapso de modos inerente ao treinamento de amostradores neurais baseados em difusão para distribuições complexas e multimodais. Ao decompor a otimização global em uma sequência de passos proximais, o método promove uma exploração completa através dos modos e estabiliza a otimização em espaços de alta dimensão.

Os autores observam duas limitações principais:

1. A investigação foca em uma única instanciação do PDNS (Proximal WDCE), deixando espaço para explorar outros objetivos proximais ou priors.
2. Os experimentos são conduzidos em benchmarks controlados em vez de aplicações reais em grande escala.

O trabalho sugere direções futuras incluem explorar problemas SOC com dinâmicas de referência não sem memória, estender o framework para espaços de estados mais gerais e desenvolver outras instanciamações do objetivo proximal.