Globally Solving Unbalanced Optimal Transport and Density Control for Gaussian Distributions

Resumen Técnico: Resolución Global del Transporte Óptimo Desbalanceado y Control de Densidad para Distribuciones Gaussianas

Formulación del Problema
El artículo aborda dos problemas interconectados: Transporte Óptimo Desbalanceado (UOT) y su extensión dinámica desde la teoría de control, Control de Densidad Desbalanceado (UDC).

1. UOT Estático: Dadas dos medidas de referencia gaussianas, $\alpha = c_\alpha \mathcal{N}(m_\alpha, \Sigma_\alpha)$ y $\beta = c_\beta \mathcal{N}(m_\beta, \Sigma_\beta)$, el objetivo es encontrar un plan de transporte $\pi$ que minimice la suma del costo de transporte cuadrático y las penalizaciones de divergencia de Kullback–Leibler (KL) sobre las marginales respecto a las referencias. A diferencia del OT balanceado clásico, la masa total del plan de transporte no está fijada a priori; las penalizaciones KL permiten que las marginales se desvíen de las referencias, tratando efectivamente la creación/destrucción de masa como una variable de optimización.
2. UDC Dinámico: Los autores extienden esto a sistemas lineales de tiempo discreto ($x_{t+1} = Ax_t + Bu_t$). El objetivo es dirigir una medida de estado desde un tiempo inicial hasta un tiempo terminal. Las distribuciones inicial y terminal no son restricciones rígidas, sino que se penalizan mediante la divergencia KL contra referencias gaussianas. La función de costo incluye el esfuerzo de control cuadrático esperado y las penalizaciones KL en los extremos. La masa total de la medida de estado se conserva mediante la dinámica, pero no está prescrita por las referencias.

Metodología y Reducción Gaussiana
La contribución metodológica central es la demostración de que estos problemas variacionales de dimensión infinita admiten reducciones exactas a dimensión finita cuando las medidas de referencia son gaussianas.

• Caso Estático (UOT): Los autores demuestran que para cualquier plan de transporte factible, existe un plan de transporte gaussiano con el mismo costo (o menor). Específicamente, la solución óptima se caracteriza por marginales gaussianas y un mapa de acoplamiento afín. Esto reduce el problema a optimizar sobre la masa ($c$), las medias ($m_1, m_2$) y las covarianzas ($\Sigma_1, \Sigma_2$).

  • La optimización se desacopla: Primero, se resuelve un problema convexo sobre los momentos (medias y covarianzas). Segundo, la masa óptima se calcula mediante una actualización escalar de forma cerrada derivada del valor óptimo del problema de momentos.
  • Este enfoque también se extiende al UOT Entrópico (EUOT), donde se incluye una penalización KL adicional sobre el plan conjunto respecto a la medida producto.

• Caso Dinámico (UDC): Los autores establecen que cualquier solución factible (que involucre medidas generales y políticas de control) puede ser reemplazada, sin pérdida de optimalidad, por una medida inicial gaussiana y una política de retroalimentación afín-gaussiana.

  • La política de control toma la forma $u_t = K_t(x_t - m_t) + v_t + w_t$, donde $w_t$ es ruido gaussiano.
  • Utilizando técnicas estándar de elevación para la dirección de covarianza ($S_t = K_t \Sigma_t$, $Y_t = K_t \Sigma_t K_t^\top + \Sigma^u_t$), las restricciones bilineales no convexas se transforman en un Programa Semidefinido (SDP).
  • Similar al caso estático, el problema se separa en un SDP convexo sobre los momentos y parámetros de control (para una masa fija) y una optimización escalar de forma cerrada para la masa óptima.
  • El artículo también establece que si la secuencia óptima de covarianzas del estado permanece definida positiva, la política óptima es determinista (es decir, la covarianza del ruido de control $\Sigma^u_t$ se anula).

• UDC de Máxima Entropía (MaxEnt UDC): El marco se extiende aún más para incluir una recompensa de entropía en la política de control. Los autores muestran que la solución óptima permanece dentro de la clase gaussiana y proporcionan una formulación convexa relajada (mediante el complemento de Schur) que se demuestra que es ajustada, garantizando la optimalidad global.

Resultados Clave y Algoritmos

• Existencia: El artículo demuestra la existencia de minimizadores globales tanto para los problemas UOT como UDC reducidos.
• Algoritmos:
  • Algoritmo 1: Resuelve el problema estático UOT resolviendo primero una optimización convexa de momentos y luego aplicando una actualización de masa de forma cerrada.
  • Algoritmo 2: Resuelve el problema UDC resolviendo un SDP para la trayectoria y los parámetros de control de masa fija, seguido de una actualización de masa de forma cerrada.
• Validación Numérica: Las simulaciones demuestran que a medida que aumenta el parámetro de penalización $\gamma$, las marginales óptimas convergen a las medidas de referencia, recuperando el comportamiento del OT balanceado estándar. En el entorno dinámico, el método dirige con éxito las medidas de estado equilibrando la coincidencia de extremos y el esfuerzo de control, con la capacidad de ajustar la masa transportada.

Significado y Afirmaciones
El artículo afirma proporcionar métodos de solución óptima global para UOT y UDC gaussianos, superando enfoques heurísticos o aproximados. El significado principal radica en:

1. Reducción Exacta: Demostrar que, a pesar de la naturaleza de dimensión infinita de los problemas originales, las soluciones óptimas para referencias gaussianas están estrictamente confinadas a la clase gaussiana, lo que permite reformulaciones exactas de dimensión finita.
2. Solubilidad Global: Al separar la optimización de masa de la optimización de momentos y utilizar relajaciones convexas (SDP), los autores proporcionan algoritmos que garantizan la optimalidad global en lugar de mínimos locales.
3. Marco Unificado: El trabajo unifica el transporte estático y el control dinámico bajo un único paradigma "desbalanceado", donde la masa es una variable de decisión en lugar de una restricción fija. Esto es particularmente relevante para aplicaciones donde las distribuciones de los extremos son perfiles de referencia en lugar de restricciones rígidas, y donde la masa total puede variar debido a la incertidumbre o la observabilidad parcial.

Los autores señalan que, aunque los problemas reducidos no son conjuntamente convexos en su forma cruda, las técnicas de separación y elevación propuestas los hacen tratables. También identifican direcciones futuras, como aplicar el marco al análisis de datos de secuenciación de ARN de células individuales y establecer conexiones con el problema del Puente de Schrödinger desbalanceado, pero no afirman haber resuelto estas aplicaciones específicas dentro de este trabajo.