Resumen Técnico: Un Enfoque Riemanniano para el Transporte Óptimo de Bajo Rango

Planteamiento del Problema

El Transporte Óptimo (OT) clásico proporciona una métrica rigurosa para comparar distribuciones de probabilidad, pero sufre un cuello de botella computacional cuadrático ($O(mn)$) debido a la necesidad de calcular matrices de acoplamiento densas. Aunque la regularización entrópica permite el uso del algoritmo de Sinkhorn, esta retiene dicha escala e introduce sesgo. Para mitigar esto, el Transporte Óptimo de Bajo Rango (LOT) restringe el plan de transporte a una estructura de baja dimensión, factorizando la matriz de acoplamiento $\Gamma$ en sub-acoplamientos $(U, V)$ que comparten una marginal latente $g$.

A pesar de la eficiencia estructural de las formulaciones de bajo rango, los solvers existentes (p. ej., LOT, LR-GW, UB-LOT) dependen fuertemente de actualizaciones de descenso de espejo de primer orden junto con bucles internos iterativos (proyecciones de Dykstra o Bregman). Estos enfoques enfrentan varias limitaciones:

1. Sensibilidad a los Hiperparámetros: Requieren un ajuste cuidadoso de los tamaños de paso y de los parámetros de regularización entrópica.
2. Curvatura Ignorada: No logran explotar la curvatura local del paisaje de optimización, lo que conduce a una convergencia más lenta.
3. Sobrecarga Computacional: Los bucles internos de iteración para la proyección y retracción aumentan la complejidad por iteración, particularmente en entornos desbalanceados.

Metodología

Los autores proponen un marco geométrico riemanniano unificado que modela el conjunto de acoplamientos factorizados positivos de rango-$r$ como subvariedades suaves embebidas dentro del ortante positivo. Al equipar estas variedades con la métrica de producto Fisher-Rao, los autores derivan formulaciones tratables para todos los ingredientes geométricos necesarios, permitiendo una optimización robusta de primer y segundo orden.

1. Formulación de la Variedad

El marco define dos variedades distintas basadas en las restricciones marginales:

• Variedad Balanceada ($\mathcal{M}_{bal}$): Impone la reconstrucción estricta de las marginales fuente y objetivo ($U\mathbf{1}=a, V\mathbf{1}=b$) y consistencia latente ($U^\top\mathbf{1}=V^\top\mathbf{1}$).
• Variedad Desbalanceada ($\mathcal{M}_{ub}$): Relaja las restricciones marginales externas, requiriendo solo consistencia latente ($U^\top\mathbf{1}=V^\top\mathbf{1}$), penalizando las desviaciones marginales mediante la divergencia KL en el objetivo.

2. Ingredientes Geométricos

Bajo la métrica Fisher-Rao, los autores derivan expresiones explícitas para:

• Proyectores Ortogonales:
  • Para $\mathcal{M}_{bal}$, la proyección implica resolver un sistema KKT (reducido a un complemento de Schur) mediante Gradiente Conjugado (CG).
  • Para $\mathcal{M}_{ub}$, la proyección se simplifica a una expresión de forma cerrada, eliminando la necesidad de solvers iterativos.
• Retracciones:
  • La retracción mapea vectores tangentes de vuelta a la variedad mediante un paso de espejo exponencial consistente con Fisher seguido de una proyección KL.
  • Para $\mathcal{M}_{bal}$, esto requiere proyecciones Bregman cíclicas (iterativas).
  • Para $\mathcal{M}_{ub}$, la proyección KL colapsa en un escalamiento de columnas de forma cerrada, eliminando los bucles internos por completo.
• Gradiente y Hessiano Riemanniano:
  • El marco proporciona una formulación agnóstica al costo donde el gradiente riemanniano y los productos Hessiano-vector dependen únicamente de las derivadas euclidianas de la función de costo específica (OT Lineal, Gromov-Wasserstein, etc.) y la geometría de la variedad.
  • Esto permite el uso de solvers de segundo orden (Región de Confianza, Gradiente Conjugado) que utilizan la información de curvatura.

3. Certificado de Suficiencia de Rango

El artículo introduce un certificado para la verificación de la optimalidad global. Al calcular un "gradiente reducido" (matriz de holgura dual) en un punto crítico de rango-$r$, el método puede verificar si la solución es globalmente óptima para el problema de rango completo. Si el rango es insuficiente, el certificado identifica una dirección para el aumento de rango.

Contribuciones Clave

1. Formulación Riemanniana Unificada: Formaliza el OT de bajo rango balanceado y desbalanceado como variedades suaves bajo la métrica Fisher-Rao, proporcionando un único marco geométrico para el OT lineal, Gromov-Wasserstein (GW) y sus variantes desbalanceadas.
2. Operaciones Desbalanceadas de Forma Cerrada: Demuestra que para el OT desbalanceado, los proyectores del espacio tangente y las retracciones se reducen a expresiones de forma cerrada. Esto elimina los bucles iterativos de Dykstra/Bregman requeridos por los métodos de vanguardia (p. ej., UB-LOT), reduciendo la complejidad por iteración a $O(rN)$.
3. Solvers de Segundo Orden: Permite el uso de solvers conscientes de la curvatura (CG, Región de Confianza) que logran una convergencia más rápida que el descenso de espejo de primer orden, sin requerir regularización entrópica ni ajuste de tamaño de paso.
4. Verificación de Optimalidad Global: Proporciona un certificado computable para verificar si una solución de bajo rango es globalmente óptima o para guiar estrategias de aumento de rango.

Resultados Experimentales

Experimentos extensos a través de entornos balanceados y desbalanceados (OT Lineal, GW, Fused GW) con tamaños de dataset de hasta $N=50,000$ demuestran:

• Velocidad de Convergencia: Los solvers riemannianos convergen en un orden de magnitud menos iteraciones externas en comparación con las líneas base de descenso de espejo.
• Tiempo de Ejecución (Wall-Clock Time): La eliminación de los bucles internos en el caso desbalanceado conduce a aceleraciones significativas. Por ejemplo, en el OT desbalanceado con valores atípicos (outliers), los solvers riemannianos convergen en segundos, mientras que UB-LOT agota el tiempo o diverge.
• Desempeño: Los métodos logran valores de objetivo superiores o comparables sin el sesgo entrópico inherente a las líneas base regularizadas.
• Ventaja de Segundo Orden: En problemas no convexos (GW, Fused GW), los métodos de Región de Confianza aprovechan la información del Hessiano para escapar de mínimos locales poco profundos, superando a los métodos de primer orden.

Significancia

El artículo sostiene que, al cambiar de un paradigma de descenso de espejo de primer orden a una perspectiva geométrica riemanniana, los solvers de OT de bajo rango pueden superar las limitaciones de mal condicionamiento y sensibilidad a los hiperparámetros. La principal significancia radica en la simplificación estructural de las operaciones desbalanceadas, que transforma un proceso iterativo computacionalmente costoso en una evaluación de forma cerrada. Esto permite una optimización sin regularización y sin necesidad de ajuste que escala linealmente con el tamaño del dataset, proporcionando al mismo tiempo garantías rigurosas de convergencia de segundo orden y verificación de optimalidad global. El marco se presenta como un enfoque general y agnóstico al costo, aplicable a un amplio espectro de problemas de transporte óptimo.