Sample complexity of unbalanced entropic OT

Resumen Técnico: Complejidad de Muestreo de OT Entrópico Desbalanceado

Planteamiento del Problema
El Transporte Óptimo (OT) proporciona un marco geométrico para comparar medidas de probabilidad, pero el OT entrópico balanceado clásico enfrenta dos limitaciones significativas en aplicaciones de aprendizaje automático: (1) la rigidez de las restricciones de conservación de masa estricta, que no es adecuada para datos con masa faltante o creada, y (2) una complejidad de muestreo de alta dimensión desfavorable (la "maldición de la dimensionalidad"), donde las tasas de convergencia empírica se deterioran con la dimensión ambiental.

Aunque se han desarrollado la regularización entrópica y las relajaciones desbalanceadas (usando $\phi$-divergencias) para abordar estos problemas individualmente, las garantías estadísticas existentes para el OT entrópico desbalanceado se limitan principalmente al valor de transporte escalar (la función objetivo). Existe una falta de límites de muestras finitas de alta probabilidad directamente sobre el acoplamiento óptimo (plan de transporte). Esta brecha es crítica porque, en aplicaciones como el modelado generativo o la inferencia de trayectorias, el acoplamiento mismo es el objeto de interés para la predicción e interpolación, no solo el valor del costo.

Metodología
El artículo desarrolla un análisis estadístico del OT entrópico desbalanceado centráquéndose en las propiedades geométricas del problema dual. La innovación metodológica central es la introducción de una formulación de envolvente invariante a la traslación para manejar la pérdida de invariancia de calibre inherente a los entornos desbalanceados.

1. Formulación Dual y Construcción de la Envolvente:
   A diferencia del OT balanceado, donde el objetivo dual es invariante bajo traslaciones aditivas de los potenciales de Kantorovich ($f \to f+\lambda, g \to g-\lambda$), el OT desbalanceado rompe esta simetría debido a que las penalizaciones de divergencia marginal dependen de los niveles absolutos de los potenciales. Los autores definen una envolvente invariante a la traslación $T_{\alpha, \beta}(f, g) := \inf_{\lambda \in \mathbb{R}} K_{\alpha, \beta}(f+\lambda, g-\lambda)$. Esta construcción separa explícitamente los grados de libertad intrínsecos del dual de la dirección de traslación espuria.

2. Análisis Geométrico (Compacidad y Convexidad Fuerte):
   Los autores demuestran que, bajo supuestos leves (espacios compactos, costo Lipschitz), esta funcional de envolvente admite un dominio de minimización anclado compacto ($P_{\alpha, \beta}^0$) donde los potenciales se fijan para tener media cero con respecto a la medida de origen. Crucialmente, en este dominio, la envolvente es uniformemente convexamente fuerte con respecto a una norma $L^2$ ponderada. Esto recupera la estructura geométrica favorable perdida en el dual desbalanceado estándar.

3. Control del Parámetro de Traslación:
   Si bien la envolvente elimina la necesidad de un ajuste de calibre arbitrario, introduce una optimización sobre un parámetro de traslación $\lambda$. El artículo establece que, bajo una condición de dominio leve sobre los generadores de divergencia (Supuesto A2), el $\lambda$ óptimo se confina a un intervalo uniformemente acotado $[-R, R]$. Esto permite límites $L^\infty$ uniformes sobre los potenciales duales.

4. Estabilidad Estadística:
   Aprovechando la convexidad fuerte de la envolvente anclada y la acotación del parámetro de traslación, los autores derivan estimaciones de estabilidad. Demuestran que los potenciales duales empíricos convergen uniformemente a los potenciales de la población. Estas tasas de convergencia dual se transfieren luego al plan de transporte primal mediante la relación de escalado exponencial ($d\pi/d(\alpha \otimes \beta) \propto \exp((f \oplus g - c)/\eta)$).

Contribuciones Clave

• Formulación Dual Invariante a la Traslación: El artículo introduce una formulación dual rigurosa que restaura la convexidad fuerte para el OT entrópico desbalanceado al cocientar la dirección de traslación, una propiedad utilizada anteriormente solo para la estabilidad algorítmica en las iteraciones de Sinkhorn.
• Complejidad de Muestreo a Nivel de Plan: La contribución principal es la derivación de límites de muestras finitas de alta probabilidad para el acoplamiento óptimo (plan de transporte) en el entorno desbalanceado, en lugar de solo para el valor del objetivo.
• Aplicabilidad a Divergencias Generales: Los resultados se aplican a una amplia clase de $\phi$-divergencias (incluyendo KL, $\chi^2$, Hellinger, Jensen–Shannon y $\alpha$-divergencias) sin requerir pruebas específicas de cada divergencia, apoyándose en argumentos generales de convexidad y compacidad.
• Constantes Explícitas: El artículo proporciona constantes explícitas para los límites sobre el parámetro de traslación y la resultante complejidad de muestreo, aclarando la dependencia del parámetro de entropía $\eta$, la función de costo y las propiedades de la divergencia.

Resultados Principales
El artículo establece que para medidas empíricas $\alpha_n, \beta_n$ derivadas de $n$ muestras, la desviación entre el plan óptimo empírico $\pi_n$ y el plan de la población $\pi^\star$ satisface un límite de alta probabilidad de la forma:
$$|\langle h, \pi^\star - \pi_n \rangle| \lesssim C \cdot e^{C(M + \|c\|_\infty)/\eta} \|h\|_\infty \sqrt{\frac{m_\alpha m_\beta (\log n + t)}{n}}$$
donde $M$ es una cota uniforme sobre los potenciales duales, $h$ es una función de prueba, y las constantes dependen de la geometría del problema y los parámetros de la divergencia. Notablemente, si la distribución de muestreo coincide con el producto de las marginales, el factor $\log n$ puede eliminarse.

Significancia y Reivindicaciones
El artículo sostiene que estos resultados aclaran la necesidad práctica de la regularización en aplicaciones de aprendizaje automático de OT. Específicamente:

• Mitigación de la Dimensionalidad: La regularización suaviza la maldición de la dimensionalidad, reduciendo el número de muestras requeridas para una estimación de transporte estable.
• Compatibilidad Algorítmica: Los estimadores resultantes siguen siendo compatibles con solvers escalables de tipo Sinkhorn, cerrando la brecha entre la teoría estadística y la práctica computacional.
• Estabilidad para Problemas Inversos: Las estimaciones de estabilidad directa proporcionadas son esenciales para los análisis complementarios de problemas inversos (por ejemplo, recuperar costos a partir de muestras), como se señala en el trabajo relacionado [1].

Los autores enfatizan que su análisis se basa en "supuestos muy leves y mayormente agnósticos a la divergencia", lo que hace que el marco teórico sea robusto y aplicable a los modelos estándar utilizados en la práctica. No pretenden resolver la maldición de la dimensionalidad por completo, sino demostrar cómo el OT desbalanceado entrópico proporciona una alternativa estadísticamente tratable a las distancias de Wasserstein no regularizadas en entornos de alta dimensión.