Proximal Diffusion Neural Sampler

Resumen Técnico: Muestreador Neuronal de Difusión Proximal (PDNS)

Enunciado del Problema
Aprender un muestreador neuronal basado en difusión para extraer muestras de una distribución objetivo no normalizada (por ejemplo, una distribución de Boltzmann $\pi(x) \propto e^{-\beta E(x)}$) es una tarea fundamental en estadística computacional, inferencia bayesiana y física estadística. Este problema se formula a menudo como un problema de Control Óptico Estocástico (SOC) sobre medidas de trayectoria, donde una red neuronal parametriza un proceso de control para transportar una distribución de referencia simple hacia la objetivo.

Sin embargo, el entrenamiento de estos muestreadores enfrenta desafíos significativos cuando la distribución objetivo es multimodal con barreras de energía altas que separan los modos. En tales escenarios, los métodos estándar basados en SOC (incluyendo variantes de Entropía Cruzada y Entropía Cruzada de Desruido Ponderada) a menudo sufren de colapso de modos. Al inicio del entrenamiento, el modelo captura solo unos pocos modos; la función de pérdida global refuerza entonces estos modos capturados basándose en las trayectorias limitadas generadas, lo que provoca que la masa de probabilidad se concentre en un subconjunto de modos ignorando a los demás. Esto conduce a una optimización inestable y a una cobertura deficiente de la distribución objetivo.

Metodología
Los autores proponen el Muestreador Neuronal de Difusión Proximal (PDNS), un marco unificado diseñado para abordar el colapso de modos reformulando el proceso de entrenamiento como una secuencia de sub-problemas locales en lugar de una única minimización global.

1. Marco Unificado: El artículo unifica primero los muestreadores de difusión continuos y discretos bajo un único marco SOC. Define la medida de trayectoria óptima $P^*$ como la solución a la minimización de un objetivo de entropía relativa sujeto a una medida de trayectoria de referencia $P_{ref}$ que satisface una condición "sin memoria" (independencia entre las distribuciones inicial y terminal).
2. Método de Punto Proximal: Para mitigar el colapso de modos, PDNS aplica el método de punto proximal al espacio de medidas de trayectoria. En lugar de optimizar directamente el objetivo global, el algoritmo resuelve iterativamente un sub-problema local en el paso $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)$$
   Aquí, $\eta_k$ es un tamaño de paso proximal, y el término $KL(P_\theta \| P_{\theta_{k-1}})$ actúa como un regularizador que mantiene la iteración actual cerca de la anterior.
3. Propiedades Teóricas: Los autores demuestran que la solución óptima al $k$-ésimo paso proximal es una interpolación geométrica entre la iteración anterior y la medida de trayectoria óptima. Específicamente, la distribución terminal del paso $k$, $P_k$, es una interpolación geométrica entre la distribución terminal de referencia $\nu$ y la objetivo $\pi$. A medida que $k \to \infty$, la secuencia de medidas de trayectoria converge a la objetivo $P^*$. Este procedimiento escalonado asegura que el modelo explore el espacio de estados gradualmente, evitando un compromiso prematuro con un único modo.
4. Instanciación Práctica (Proximal WDCE): Para lograr eficiencia, los autores instancian PDNS utilizando un objetivo de Entropía Cruzada de Desruido Ponderada Proximal (Proximal WDCE). Este enfoque evita almacenar trayectorias completas mediante el uso de pesos de importancia acumulados en línea.
   • Caso Continuo: El objetivo implica igualar una puntuación de desruido (o puente) con un término de peso derivado del tamaño de paso proximal $\eta_k$ y la función de energía.
   • Caso Discreto: El objetivo adapta la difusión discreta enmascarada, utilizando un esquema de ponderación similar basado en la derivada de Radon-Nikodym de las medidas de trayectoria.
5. Diseño del Programador: El artículo discute estrategias para seleccionar el tamaño de paso $\eta_k$ (o su equivalente $\lambda_k$). Un $\eta_k$ más pequeño (regularización más fuerte) preserva la cobertura de modos pero ralentiza la convergencia, mientras que un $\eta_k$ más grande arriesga el colapso de modos. Los autores proponen tanto programas predefinidos ($\lambda_k$ decreciendo linealmente) como una estrategia adaptativa que ajusta $\eta_k$ basándose en una restricción de región de confianza sobre la divergencia KL entre el modelo actual y la siguiente objetivo.

Contribuciones Clave

1. Marco Unificado: El artículo integra muestreadores basados en SOC continuos y discretos en una única formulación de medida de trayectoria.
2. Muestreador Neuronal de Difusión Proximal (PDNS): Introduce un marco novedoso que aprovecha iteraciones de punto proximal en el espacio de trayectorias para estabilizar el aprendizaje y mitigar el colapso de modos en configuraciones multimodales.
3. Proximal WDCE: Deriva instanciaciones concretas y computacionalmente eficientes de PDNS para dominios continuos y discretos utilizando pérdidas de entropía cruzada de desruido ponderado proximal.
4. Estrategias de Programador: Los autores delinean estrategias fundamentadas para seleccionar tamaños de paso proximales, incluyendo un programador adaptativo que equilibra la velocidad de convergencia y la estabilidad.

Resultados Experimentales
Los autores evalúan PDNS en una amplia gama de puntos de referencia continuos y discretos:

• Tareas Continuas:
  • Funciones de Energía Sintéticas: PDNS supera a bases sólidas (incluyendo SMC, DDS, PIS y AS) en puntos de referencia multimodales como el potencial de 32 pozos (MW54), Funnel y mezclas gaussianas de alta dimensión (GMM40). Logra resultados del estado del arte en la mezcla de distribuciones t de Student (MoS) y potenciales basados en partículas (sistemas de Lennard-Jones LJ-13 y LJ-55).
  • Dinámica Molecular: En el sistema de Dipéptido de Alanina (AD), PDNS logra divergencias KL comparables al método ASBS del estado del arte y captura con precisión las estructuras conformacionales en los gráficos de Ramachandran.
• Tareas Discretas:
  • Física Estadística: PDNS supera significativamente a LEAPS y Metropolis-Hastings (MH) en el muestreo de modelos de Ising y Potts a bajas temperaturas (cerca o por debajo de los puntos críticos), donde la separación de modos es severa. Logra errores menores en magnetización y correlaciones de 2 puntos.
  • Optimización Combinatoria: Aplicado al problema del Corte Máximo en grafos Barabási–Albert y Erdős–Rényi, PDNS demuestra la capacidad de alcanzar calidades de solución comparables a solucionadores de verdad fundamental (Gurobi).

Significado y Afirmaciones
El artículo afirma que PDNS proporciona una solución robusta al problema de colapso de modos inherente al entrenamiento de muestreadores neuronales basados en difusión para distribuciones complejas y multimodales. Al descomponer la optimización global en una secuencia de pasos proximales, el método promueve una exploración exhaustiva a través de los modos y estabiliza la optimización en espacios de alta dimensión.

Los autores señalan dos limitaciones principales:

1. La investigación se centra en una única instanciación de PDNS (Proximal WDCE), dejando margen para explorar otros objetivos o priores proximales.
2. Los experimentos se realizan en puntos de referencia controlados en lugar de aplicaciones del mundo real a gran escala.

El trabajo sugiere direcciones futuras que incluyen explorar problemas SOC con dinámicas de referencia no sin memoria, extender el marco a espacios de estados más generales y desarrollar otras instanciaciones del objetivo proximal.