Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization

Resumen Técnico: Diferencia de Programación Convexa en el Espacio de Wasserstein con Aplicaciones a la Optimización de MMD

Planteamiento del Problema
La optimización de funcionales sobre el espacio de medidas de probabilidad es central para tareas de aprendizaje automático que van desde la inferencia variacional hasta el modelado generativo. Un enfoque estándar consiste en realizar la optimización directamente sobre el espacio de Wasserstein $(P_2(\mathbb{R}^d), W_2)$ mediante la discretización de flujos de gradiente de Wasserstein. Si bien esto es efectivo para funcionales convexos (p. ej., divergencia KL, divergencias f), muchos objetivos prácticamente relevantes son no convexos a lo largo de las geodésicas de Wasserstein. Ejemplos notables incluyen la distancia de Wasserstein al cuadrado, la distancia de Sliced-Wasserstein y la Discrepancia de la Media Máxima (MMD). Los métodos de primer orden estándar como el Descenso de Gradiente de Wasserstein (WGD) tienen dificultades en estos entornos no convexos, requiriendo a menudo una inyección delicada de ruido o la optimización en espacios alternativos (p. ej., espacios Wasserstein-MMD o Wasserstein-Fisher-Rao) que alteran los pesos de la distribución o mezclan pasos implícitos.

Metodología
Los autores proponen un nuevo algoritmo basado en partículas diseñado para funcionales no convexos que admiten una descomposición de Diferencia de Convexos (DC) dentro del espacio de Wasserstein. La metodología central consiste en elevar el clásico Procedimiento Convexa-Cóncavo (CCCP) desde los espacios euclidianos al espacio de Wasserstein, resultando en el Procedimiento Convexo-Cóncavo de Wasserstein (WCCCP).

1. Descomposición DC en el Espacio de Wasserstein: El funcional objetivo $F$ se descompone como $F(\mu) = F_+(\mu) - F_-(\mu)$, donde tanto $F_+$ como $F_-$ son funcionales totalmente convexos en $P_2(\mathbb{R}^d)$. A diferencia de trabajos previos que restringían la parte cóncava a potenciales (funcionales lineales), este enfoque permite que $F_-$ sea cualquier funcional diferenciable.
2. Esquema Algorítmico: En cada iteración $k$, dada una medida $\mu_k$, el algoritmo resuelve un problema de mayoración-minimización. Se linealiza la parte cóncava $F_-$ alrededor de $\mu_k$ para construir una cota superior y se minimiza el surrogado convexo resultante sobre el espacio de mapas de transporte $T \in L^2(\mu_k)$:
   $$T_{k+1} = \arg\min_{T \in L^2(\mu_k)} \left( F_+(T\#\mu_k) - \langle \nabla_{W_2}F_-(\mu_k), T - \text{Id} \rangle_{L^2(\mu_k)} \right)$$
   La siguiente iteración se define entonces como $\mu_{k+1} = (T_{k+1})\#\mu_k$.
3. Implementación: La minimización interna se realiza mediante descenso de gradiente en $L^2(\mu_k)$. Para distribuciones empíricas (partículas), esto se traduce en la actualización de las posiciones de las partículas. Los autores también establecen conexiones entre WCCCP y el Descenso de Espejo/Bregman Proximal en el espacio de Wasserstein.
4. Especificidades de MMD: Para la Discrepancia de la Media Máxima, los autores proponen descomposiciones DC explícitas dividiendo el propio núcleo $k(x,y) = \psi(x-y)$ en componentes convexos $\psi = \psi_+ - \psi_-$. Proporcionan descomposiciones específicas para núcleos invariantes por traslación, incluyendo núcleos Gaussianos y de Riesz, utilizando la descomposición de Jordan o división algebraica (p. ej., $\cosh/\sinh$).

Contribuciones Clave

• Algoritmo WCCCP: Introducción del Procedimiento Convexo-Cóncavo de Wasserstein, extendiendo la optimización DC al espacio de Wasserstein sin restringir el componente cóncavo a potenciales lineales.
• Convergencia Teórica:
  • Prueba de que el valor del objetivo $F(\mu_k)$ es no creciente a lo largo de las iteraciones.
  • Establecimiento de resultados de estacionariedad casi total. Bajo supuestos de suavidad y convexidad fuerte (convexidad total) de los componentes, los autores demuestran tasas de convergencia sublineales para la distancia al cuadrado entre iteraciones y la norma del gradiente de Wasserstein. Específicamente, $\min_{k} \|\nabla_{W_2}F(\mu_k)\|^2 = O(1/K)$.
  • Demostración de que el esquema es equivalente al Descenso Proximal de Bregman y al Descenso de Espejo bajo condiciones específicas.
• Descomposiciones DC Explícitas para MMD: El artículo proporciona una forma sistemática de descomponer funcionales MMD con núcleos invariantes por traslación en formas DC, satisfaciendo las condiciones necesarias de convexidad y suavidad para las garantías de convergencia.
• Validación Empírica: Los experimentos con la Distancia de Energía y MMD con núcleos Gaussianos demuestran que las descomposiciones DC bien elegidas producen una convergencia más rápida y estable en comparación con el WGD estándar y el método de Gradiente Proximal de Wasserstein (Forward-Backward).

Resultados

• Distancia de Energía: En datos sintéticos (formas de espiral y gato) e imágenes de CIFAR10, WCCCP converge significativamente más rápido que WGD y el esquema Forward-Backward (FB). El algoritmo logra minimizar la Distancia de Energía para igualar distribuciones objetivo complejas.
• MMD con Núcleo Gaussiano: Minimizar la MMD con un núcleo Gaussiano es notoriamente sensible al ancho de banda y a la inicialización. Los autores muestran que el WGD y el esquema FB (usando una descomposición estándar) a menudo se quedan atrapados en mínimos locales donde las partículas se alejan de la distribución objetivo. En contraste, WCCCP utilizando descomposiciones de núcleo específicas (Jordan o $\cosh/\sinh$) logra una mejor convergencia. Se observa que la descomposición de Jordan es más suave y menos propensa a la inestabilidad numérica, mientras que la descomposición $\cosh/\sinh$ asegura que todas las partículas eventualmente converjan, aunque a un ritmo más lento.
• Comparación: Los resultados resaltan que el rendimiento de WCCCP depende altamente de la elección de la descomposición DC. Una descomposición "bien elegida" es crítica para superar al descenso de gradiente estándar.

Significancia y Afirmaciones
El artículo afirma proporcionar un marco robusto para optimizar funciales no convexos en el espacio de Wasserstein, abordando una brecha donde los métodos de gradiente estándar fallan debido a la no convexidad. Al elevar el CCCP a este entorno, los autores ofrecen una alternativa teóricamente fundamentada que evita los problemas de mezcla de pesos de los enfoques anteriores específicos de MMD. El trabajo enfatiza que, si bien el algoritmo es general, su éxito práctico depende fuertemente de la descomposición DC específica del núcleo. Los autores señalan modestamente que se requiere trabajo futuro para comprender más profundamente el impacto de diferentes descomposiciones y para diseñar estrategias adaptativas automáticas para su selección. Además, el análisis teórico actualmente asume soluciones exactas a los problemas de optimización interna, y el considerar solucionadores inexactos se identifica como una dirección necesaria para la investigación futura.