Resumen Técnico: Transporte Óptimo Neuronal Riemaniano Entrópico

1. Enunciado del Problema

Muchas aplicaciones de aprendizaje automático involucran datos soportados en espacios curvos (variedades riemannianas) como esferas ($S^2$), grupos de rotación ($SO(3)$), poses rígidas ($SE(3)$) y matrices simétricas definidas positivas ($SPD$). En estos contextos, las aproximaciones euclidianas estándar distorsionan distancias, promedios y los problemas de Transporte Óptimo (OT) resultantes.

Los enfoques existentes enfrentan una compensación:

• Los métodos de OT en variedades a menudo persiguen mapas de transporte amortizados y fuera de muestra, pero sufren de cuellos de botella computacionales, requiriendo frecuentemente optimizaciones internas iterativas para cada nueva instancia.
• La Regularización Entrópica (por ejemplo, iteraciones de Sinkhorn) hace que el OT discreto sea escalable y numéricamente estable, pero no proporciona intrínsecamente un modelo amortizado; cada nuevo par de distribuciones típicamente requiere resolver un nuevo problema de optimización.

El artículo aborda la brecha de combinar OT geométrico intrínseco con evaluación amortizada fuera de muestra y regularización entrópica en variedades riemannianas posiblemente no compactas.

2. Metodología: RNOT Entrópico

Los autores proponen Transporte Óptimo Neuronal Riemaniano Entrópico (Entropic RNOT), un marco unificado que aprende un modelo de transporte reutilizable y consciente de la variedad.

Formulación Central

El método se basa en la formulación semidual del OT entrópico. En lugar de aprender un mapa de transporte directamente, el modelo aprende un potencial de Schrödinger en el lado objetivo $g_\theta$.

• Parametrización: El potencial se parametriza mediante una retrotracción neuronal. Un mapa de características continuo $\phi: K_\nu \to \mathbb{R}^n$ (donde $K_\nu$ es el soporte de la distribución objetivo) mapea puntos de la variedad al espacio euclidiano. Una red neuronal euclidiana $a_\theta$ se compone con $\phi$ para formar la clase de hipótesis.
• Centrado: Dado que los potenciales de Schrödinger son identificables solo hasta una constante aditiva, el modelo utiliza una clase de retrotracción centrada $C_\nu(\phi^* \mathcal{F})$ para garantizar la unicidad.
• Optimización: El modelo se entrena maximizando el objetivo semidual $J_\varepsilon(g_\theta)$ mediante ascenso de gradiente estocástico en mini-lotes. El potencial en el lado fuente $f^\varepsilon_\theta$ se recupera mediante la transformación $c$ suave (una operación log-sum-exp) del potencial objetivo aprendido.

Sustitutos de Transporte Intrínseco

Una vez que el acoplamiento de Gibbs $\pi^\varepsilon_\theta$ es inducido por los potenciales aprendidos, el artículo extrae sustitutos de transporte deterministas adecuados para diferentes geometrías de variedad:

1. Proyecciones Barycéntricas: En variedades de Cartan–Hadamard (completas, simplemente conexas, de curvatura no positiva), las leyes condicionales definen un mapa de transporte determinista mediante el barycentro riemanniano (media de Fréchet).
2. Sustitutos Suavizados por Calor: En variedades estocásticamente completas (una clase más amplia que incluye variedades compactas, espacios euclidianos y productos como $SE(3)$), el método aplica suavizado por calor a las leyes condicionales objetivo. Esto convierte distribuciones condicionales potencialmente atómicas (de muestras finitas) en distribuciones absolutamente continuas. Luego, se deriva una predicción puntual (modo) a partir de esta densidad suavizada.

3. Contribuciones Clave

El artículo presenta tres contribuciones principales:

1. Introducción del Marco: Entropic RNOT es el primer marco neuronal intrínseco para OT entrópico estático en variedades riemannianas que combina la formulación semidual con la evaluación amortizada fuera de muestra.
2. Garantías Teóricas: Para un parámetro de regularización fijo $\varepsilon > 0$, los autores demuestran que la clase de hipótesis propuesta puede recuperar el acoplamiento óptimo entrópico en métricas probabilísticas fuertes (divergencia KL, Variación Total, convergencia débil). En consecuencia:
   • Los sustitutos barycéntricos convergen en $L^2(\mu)$ en variedades de Cartan–Hadamard.
   • Los sustitutos suavizados por calor son estables en cualquier tiempo de calor fijo $t > 0$ y son asintóticamente no sesgados cuando $t \to 0$.
   • Estas garantías se mantienen para datos con soporte compacto en variedades posiblemente no compactas.
3. Validación Empírica: El método demuestra una fuerte calidad de transporte en geometrías diversas ($S^2, SO(3), SPD(3), SE(3), H^2$), superando las líneas base euclidianas ambientales, de espacio tangente y log-euclidianas. Escala favorablemente en memoria y tiempo en comparación con el Sinkhorn discreto en variedades y logra mejoras significativas en una aplicación real de acoplamiento proteína-ligando.

4. Resultados Experimentales

Benchmarks Sintéticos

Evaluado en $S^2, SO(3), SPD(3), SE(3)$ y $H^2$ con distribuciones normales envueltas.

• Precisión: Entropic RNOT recupera consistentemente el plan de referencia de Sinkhorn discreto en variedad con mayor precisión que todas las líneas base, observándose las mayores ganancias en $SPD(3)$, $SE(3)$y$H^2$, donde la geometría intrínseca es más crítica.
• Métricas: Logra una divergencia KL del plan significativamente menor y errores geodésicos de punto final menores en comparación con los métodos de linealización euclidianos ambientales y de espacio tangente.

Escalabilidad

• Complejidad: El Sinkhorn discreto en variedad requiere una huella de memoria de $O(N^2)$ para la matriz de costos, volviéndose inviable para tamaños de soporte grandes (por ejemplo, $N=32,768$).
• Rendimiento: El tiempo de entrenamiento y el uso de memoria de Entropic RNOT permanecen constantes con respecto al tamaño de soporte $N$ (dependen únicamente del tamaño del lote). El rendimiento de inferencia escala linealmente con $N$, permitiendo el procesamiento de millones de muestras por segundo.

Aplicación Real: Acoplamiento Proteína-Ligando

El método se aplicó para refinar poses rígidas en $SE(3)$ para el acoplamiento proteína-ligando utilizando el conjunto de datos CrossDocked2020.

• Configuración: Se entrenó un único modelo en complejos agrupados para refinar las poses de acoplamiento retenidas hacia la cuenca de unión mejor clasificada del motor de acoplamiento. No se utilizaron estructuras cristalográficas durante el entrenamiento o la inferencia.
• Resultados:
  • Reducción del RMSD top-1 de 11.24 Å (sin refinamiento) a 3.47 Å.
  • Mejora de la tasa de éxito dentro de 2 Å del 10.3% al 75.9%.
  • Superó tanto a la minimización basada en física (GNINA) como al Sinkhorn discreto por instancia (que falló debido a conjuntos objetivo pequeños por complejo).

5. Significado y Limitaciones

Significado:
El artículo afirma proporcionar el primer marco neuronal intrínseco que unifica la escalabilidad de la regularización entrópica con las capacidades de generalización del OT neuronal amortizado en variedades. Ofrece una solución práctica para tareas de transporte de alta dimensión y no euclidianas donde los métodos discretos son computacionalmente prohibitivos.

Limitaciones (según lo declarado por los autores):

• Alcance Teórico: Las garantías teóricas se establecen para $\varepsilon > 0$ fijo y soportes compactos; no se aborda el régimen de regularización desvaneciente ($\varepsilon \to 0$).
• Restricciones Geométricas: Las garantías de recuperación de mapas barycéntricos requieren el entorno de Cartan–Hadamard; fuera de este, los barycentros pueden no ser únicos o ser inestables.
• Especificidades de la Aplicación: En el experimento de acoplamiento, el método actúa como un procedimiento de refinamiento/desruido para conjuntos de poses existentes en lugar de un modelo generativo de novo. Actualmente ignora el contexto del bolsillo del receptor y trata a los ligandos como cuerpos rígidos, omitiendo la flexibilidad torsional.
• Dependencias Computacionales: El rendimiento depende de una evaluación eficiente de la distancia geodésica y de cálculos estables de log-sum-exp.