Fixed-Point Neural Optimal Transport without Implicit Differentiation

Resumen Técnico: Transporte Óptimo Neuronal de Punto Fijo sin Diferenciación Implícita

Planteamiento del Problema
El Transporte Óptimo (OT) busca una aplicación entre distribuciones de probabilidad que minimice un costo de transporte prescrito. Si bien enfoques recientes de aprendizaje profundo han permitido solvers de OT escalables, estos dependen frecuentemente de optimización min–max adversaria (por ejemplo, WGANs) o arquitecturas de múltiples redes que involucran mapas de transporte y potenciales duales. Estos métodos suelen sufrir inestabilidad en el entrenamiento, sensibilidad a hiperparámetros y dificultades de convergencia, particularmente en configuraciones de alta dimensión. Los enfoques alternativos basados en flujos continuos requieren resolver ecuaciones diferenciales ordinarias (ODE), incurriendo en una sobrecarga computacional significativa, mientras que las formulaciones de Hamilton–Jacobi–Bellman (HJB) luchan con la escalabilidad en altas dimensiones. Además, muchos métodos existentes requieren redes auxiliares o procedimientos de optimización alternada para imponer la factibilidad dual.

Metodología
Los autores proponen un marco de red única para el transporte óptimo cuadrático que elimina el entrenamiento adversario y la diferenciación implícita. El núcleo del método es una formulación neuronal implícita derivada del problema dual de Kantorovich.

1. Parametrización de Potencial Único: El potencial dual $\phi$ se parametriza como el negativo de una función convexa $g_\theta$ representada por una única red neuronal ($\phi_\theta(x) = -g_\theta(x)$). Esto aprovecha el teorema de Brenier, que establece que el mapa de transporte óptimo es el gradiente de un potencial convexo.
2. Reformulación de Punto Fijo Proximal: Bajo esta parametrización, la $c$-transformada (un componente clave del objetivo dual) se reduce a un problema de optimización proximal fuertemente convexo:
   $$(-g_\theta)^c(z) = \inf_{y \in \mathbb{R}^d} \left\{ \frac{1}{2}\|y - z\|^2 + g_\theta(y) \right\}$$
   El minimizador $y^*_\theta(z)$ de este problema define el mapa de transporte inverso ($T_{\nu \to \mu}(z) = y^*_\theta(z)$), mientras que el mapa directo viene dado por $T_{\mu \to \nu}(x) = x + \nabla g_\theta(x)$.
3. Entrenamiento mediante Iteración de Punto Fijo: El objetivo de entrenamiento implica minimizar la pérdida dual negativa. Crucialmente, la minimización interna (el operador proximal) se resuelve utilizando iteraciones estándar de punto fijo (por ejemplo, pasos de descenso de gradiente) en lugar de optimización desenrollada o diferenciación implícita.
4. Cálculo de Gradientes sin Diferenciación Implícita: Una idea teórica clave es que el gradiente de la pérdida con respecto a los parámetros de la red $\theta$ no requiere retropropagación a través de las iteraciones de punto fijo. Debido a las condiciones de optimalidad de primer orden del operador proximal, los términos que involucran la derivada del punto fijo con respecto a $\theta$ ($dy^*/d\theta$) se cancelan exactamente. En consecuencia, el gradiente depende únicamente de las derivadas explícitas de $g_\theta$ evaluadas en las muestras de entrada y en las soluciones aproximadas de punto fijo. Esto permite un entrenamiento eficiente utilizando retropropagación estándar.
5. Convergencia: Los autores establecen que el Descenso de Gradiente Estocástico (SGD) converge a un punto estacionario incluso cuando el problema interno de punto fijo se resuelve de manera aproximada (dentro de una tolerancia especificada), siempre que el error de aproximación esté acotado.
6. Extensión Condicional a Clases: El marco se extiende naturalmente al OT condicional a clases condicionando el potencial $g_\theta(x, k)$ sobre las etiquetas de clase $k$, permitiendo que una única red maneje múltiples tareas de transporte simultáneamente.

Contribuciones Clave

• Reformulación de Red Única: Los autores derivan una formulación de OT cuadrático que requiere únicamente una red neuronal para recuperar tanto los mapas de transporte directos como inversos, eliminando la necesidad de configuraciones min–max adversarias o redes conjugadas auxiliares.
• Entrenamiento Libre de Diferenciación Implícita: Demuestran que el objetivo puede diferenciarse sin diferenciación implícita, confiando en la cancelación exacta de términos inducida por las condiciones de optimalidad del operador proximal.
• Garantías de Convergencia: El artículo proporciona pruebas teóricas para la convergencia del SGD a un punto estacionario bajo evaluaciones de punto fijo inexactas, abordando la realidad práctica de que el bucle interno no se resuelve con precisión de máquina.
• Marco Unificado: El método recupera simultáneamente los mapas directos e inversos y se extiende naturalmente a configuraciones condicionales a clases sin cambios arquitectónicos.

Resultados Experimentales
El método propuesto se evaluó en benchmarks gaussianos de alta dimensión, conjuntos de datos de física del mundo real (UCI) y tareas de traducción de imágenes (de MNIST a Fashion-MNIST).

• Gaussianos de Alta Dimensión: En benchmarks gaussianos con dimensiones de hasta $d=64$, el método logró un Porcentaje de Varianza No Explicada (UVP) significativamente menor en comparación con líneas base como NOT, MM-v1, MM:R, NCF y WGAN-QC. Mantuvo una precisión estable a medida que aumentaba la dimensionalidad, mientras que las líneas base mostraron degradación.
• Eficiencia Computacional: El método demostró una eficiencia computacional y de memoria favorable, requiriendo menos parámetros y menos tiempo de entrenamiento que los métodos basados en flujos continuos (FFJORD, RNODE) y otros enfoques neuronales de OT.
• Datos del Mundo Real: En diversos conjuntos de datos UCI (Power, Gas, HEPMASS, MINIBOONE) y parches de imágenes BSDS300, el método logró un rendimiento de coincidencia de distribuciones competitivo o superior (medido por la Discrepancia Media Máxima) con tamaños de modelo sustancialmente menores.
• Traducción de Imágenes: En tareas de traducción de imágenes condicionales a clases, el método logró la menor Distancia Fréchet Inception (FID) y la mayor precisión de clasificación entre las líneas base, indicando una preservación superior de la estructura de clase y el contenido semántico.
• Estudios de Ablación: Los experimentos mostraron que el método es robusto a la tolerancia de parada del solver de punto fijo; una inexactitud moderada en el bucle interno produce un compromiso favorable entre precisión y velocidad. Además, si bien imponer convexidad estricta mediante proyección de pesos estabilizó las iteraciones de punto fijo, degradó la precisión general del transporte, lo que sugiere que una parametrización no convexa (sin proyección rígida) ofrece un mejor equilibrio entre expresividad y estabilidad.

Significado y Afirmaciones
El artículo afirma que este marco ofrece una alternativa estable, eficiente y escalable a los métodos neuronales de OT existentes. Al eliminar el entrenamiento adversario y la sobrecarga computacional de los solvers de ODE o la resolución de ecuaciones HJB sobre todo el dominio, el método proporciona un mecanismo "libre de causalidad" para imponer condiciones de optimalidad. Los autores afirman que su enfoque es ampliamente aplicable, capaz de manejar problemas de alta dimensión y condicionales a clases mientras recupera mapas de transporte precisos en ambas direcciones. La obra se posiciona como un paso hacia la simplificación del paisaje de optimización del OT, alejándose de formulaciones min–max complejas hacia un problema de minimización de un solo objetivo que está fundamentado teóricamente y es eficiente en la práctica.