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

Résumé Technique : Différence de Programmation Convexe dans l'Espace de Wasserstein avec Applications à l'Optimisation de la MMD

Énoncé du Problème
L'optimisation de fonctionnelles sur l'espace des mesures de probabilité est centrale dans de nombreuses tâches d'apprentissage automatique, allant de l'inférence variationnelle à la modélisation générative. Une approche standard consiste à effectuer l'optimisation directement sur l'espace de Wasserstein $(P_2(\mathbb{R}^d), W_2)$ en discrétisant les flux de gradient de Wasserstein. Bien que cette méthode soit efficace pour les fonctionnelles convexes (ex. : divergence KL, divergences f-), de nombreux objectifs pratiquement pertinents sont non-convexes le long des géodésiques de Wasserstein. Des exemples notables incluent la distance de Wasserstein au carré, la distance de Wasserstein découpée (Sliced-Wasserstein) et la Maximum Mean Discrepancy (MMD). Les méthodes de premier ordre standard comme la Descente de Gradient de Wasserstein (WGD) peinent dans ces contextes non-convexes, nécessissant souvent une injection délicate de bruit ou une optimisation dans des espaces alternatifs (ex. : espaces Wasserstein-MMD ou Wasserstein-Fisher-Rao) qui modifient les poids des distributions ou mélangent des étapes implicites.

Méthodologie
Les auteurs proposent un nouvel algorithme basé sur les particules, adapté aux fonctionnelles non-convexes admettant une décomposition de type Différence de Convexité (DC) dans l'espace de Wasserstein. La méthodologie centrale consiste à élever la procédure classique de Convexité-Concavité (CCCP) des espaces euclidiens vers l'espace de Wasserstein, résultant en la Procédure de Convexité-Concavité de Wasserstein (WCCCP).

1. Décomposition DC dans l'Espace de Wasserstein : La fonctionnelle d'objectif $F$ est décomposée en $F(\mu) = F_+(\mu) - F_-(\mu)$, où $F_+$ et $F_-$ sont des fonctionnelles totalement convexes sur $P_2(\mathbb{R}^d)$. Contrairement aux travaux antérieurs qui restreignaient la partie concave à des potentiels d'énergie (fonctionnelles linéaires), cette approche permet à $F_-$ d'être n'importe quelle fonctionnelle différentiable.
2. Schéma Algorithmique : À chaque itération $k$, étant donné une mesure $\mu_k$, l'algorithme résout un problème de majoration-minimisation. Il linéarise la partie concave $F_-$ autour de $\mu_k$ pour construire une borne supérieure et minimise la surrograte convexe résultante sur l'espace des cartes de transport $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 nouvelle itérée est ensuite définie par $\mu_{k+1} = (T_{k+1})\#\mu_k$.
3. Implémentation : La minimisation interne est effectuée via une descente de gradient sur $L^2(\mu_k)$. Pour les distributions empiriques (particules), cela se traduit par une mise à jour des positions des particules. Les auteurs établissent également des connexions entre WCCCP et la Descente de Miroir (Mirror Descent) ou la Descente de Bregman Proximal dans l'espace de Wasserstein.
4. Spécificités de la MMD : Pour la Maximum Mean Discrepancy, les auteurs proposent des décompositions DC explicites en décomposant le noyau $k(x,y) = \psi(x-y)$ lui-même en composantes convexes $\psi = \psi_+ - \psi_-$. Ils fournissent des décompositions spécifiques pour les noyaux invariants par translation, incluant les noyaux Gaussien et de Riesz, en utilisant la décomposition de Jordan ou une séparation algébrique (ex. : $\cosh/\sinh$).

Contributions Clés

• Algorithme WCCCP : Introduction de la procédure de Convexité-Concavité de Wasserstein, étendant l'optimisation DC à l'espace de Wasserstein sans restreindre la composante concave à des potentiels linéaires.
• Convergence Théorique :
  • Preuve que la valeur de la fonctionnelle $F(\mu_k)$ est non-croissante le long des itérées.
  • Établissement de résultats de quasi-stationnarité. Sous des hypothèses de lissage et de forte convexité (convexité totale) des composantes, les auteurs prouvent des taux de convergence sous-linéaires pour la distance au carré entre les itérées et la norme du gradient de Wasserstein. Spécifiquement, $\min_{k} \|\nabla_{W_2}F(\mu_k)\|^2 = O(1/K)$.
  • Démonstration que le schéma est équivalent à la Descente de Bregman Proximal et à la Descente de Miroir sous certaines conditions.
• Décompositions DC Explicites pour la MMD : L'article fournit un moyen systématique de décomposer les fonctionnelles MMD avec des noyaux invariants par translation en formes DC, satisfaisant les conditions de convexité et de lissage nécessaires pour les garanties de convergence.
• Validation Empirique : Les expériences sur la Distance d'Énergie et la MMD avec des noyaux Gaussiens démontrent que des décompositions DC bien choisies produisent une convergence plus rapide et plus stable que la WGD standard et la méthode de Gradient Proximal de Wasserstein (Forward-Backward).

Résultats

• Distance d'Énergie : Sur des données synthétiques (formes de spirales et de chats) et des images CIFAR10, WCCCP converge significativement plus vite que WGD et le schéma Forward-Backward (FB). L'algorithme minimise avec succès la Distance d'Énergie pour correspondre à des distributions cibles complexes.
• MMD avec Noyau Gaussien : Minimiser la MMD avec un noyau Gaussien est notoirement sensible à la largeur de bande (bandwidth) et à l'initialisation. Les auteurs montrent que la WGD et le schéma FB (utilisant une décomposition standard) restent souvent bloqués dans des minima locaux où les particules s'éloignent de la cible. En revanche, WCCCP utilisant des décompositions de noyaux spécifiques (Jordan ou $\cosh/\sinh$) atteint une meilleure convergence. La décomposition de Jordan est notée comme étant plus lisse et moins sujette à l'instabilité numérique, tandis que la décomposition $\cosh/\sinh$ garantit que toutes les particules finissent par converger, bien qu'à un rythme plus lent.
• Comparaison : Les résultats soulignent que la performance de WCCCP dépend fortement du choix de la décomposition DC. Une décomposition « bien choisie » est cruciale pour surpasser la descente de gradient standard.

Signification et Revendications
L'article prétend fournir un cadre robuste pour optimiser des fonctionnelles non-convexes dans l'espace de Wasserstein, comblant une lacune où les méthodes de gradient standard échouent en raison de la non-convexité. En élevant la CCCP dans ce cadre, les auteurs offrent une alternative théoriquement fondée qui évite les problèmes de mélange de poids des approches précédentes spécifiques à la MMD. Le travail souligne que si l'algorithme est général, son succès pratique repose lourdement sur la décomposition DC spécifique du noyau. Les auteurs notent modestement que des travaux futurs sont nécessaires pour comprendre plus profondément l'impact des différentes décompositions et pour concevoir des stratégies automatiques et adaptatives de sélection. De plus, l'analyse théorique suppose actuellement des solutions exactes aux problèmes d'optimisation internes, et la prise en compte de solveurs inexacts est identifiée comme une direction nécessaire pour la recherche future.