Résumé Technique : Une Approche Riemannienne pour le Transport Optimal de Bas Rang

Énoncé du Problème

Le Transport Optimal (OT) classique fournit une métrique rigoureuse pour comparer des distributions de probabilité, mais souffre d'un goulot d'étranglement computationnel quadratique ($O(mn)$ dû à la nécessité de calculer des matrices de couplage denses). Bien que la régularisation entropique permette l'utilisation de l'algorithme de Sinkhorn, elle conserve cette complexité et introduit un biais. Pour atténuer cela, le Transport Optimal de Bas Rang (LOT) contraint le plan de transport à une structure de faible dimension, en factorisant la matrice de couplage $\Gamma$ en sous-couplages $(U, V)$ partageant une marge latente $g$.

Malgré l'efficacité structurelle des formulations de bas rang, les solveurs existants (par exemple, LOT, LR-GW, UB-LOT) reposent fortement sur des mises à jour de descente de miroir de premier ordre couplées à des boucles internes itératives (projections de Dykstra ou de Bregman). Ces approches font face à plusieurs limitations :

1. Sensibilité aux Hyperparamètres : Elles nécessitent un réglage minutieux des tailles de pas et des paramètres de régularisation entropique.
2. Courbure Ignorée : Elles ne parviennent pas à exploiter la courbure locale du paysage d'optimisation, ce qui conduit à une convergence plus lente.
3. Surcharge Computationnelle : Les boucles itératives internes pour la projection et la rétraction augmentent la complexité par itération, particulièrement dans les contextes non équilibrés (unbalanced).

Méthodologie

Les auteurs proposent un cadre géométrique riemannien unifié qui modélise l'ensemble des couplages factorisés de rang-$r$ comme des sous-variétés lisses enchâssées dans l'orthant positif. En dotant ces variétés de la métrique produit de Fisher-Rao, les auteurs dérivent des formulations tractables pour tous les composants géométriques nécessaires, permettant une optimisation robuste de premier et second ordre.

1. Formulation de la Variété (Manifold)

Le cadre définit deux variétés distinctes basées sur les contraintes de marges :

• Variété Équilibrée ($\mathcal{M}_{bal}$) : Impose une reconstruction stricte des marges sources et cibles ($U\mathbf{1}=a, V\mathbf{1}=b$) et une cohérence latente ($U^\top\mathbf{1}=V^\top\mathbf{1}$).
• Variété Non Équilibrée ($\mathcal{M}_{ub}$) : Relâche les contraintes de marges externes, exigeant uniquement la cohérence latente ($U^\top\mathbf{1}=V^\top\mathbf{1}$), avec les écarts de marges pénalisés via la divergence KL dans l'objectif.

2. Composants Géométriques

Sous la métrique de Fisher-Rao, les auteurs dérivent des expressions explicites pour :

• Projecteurs Orthogonaux :
  • Pour $\mathcal{M}_{bal}$, la projection implique la résolution d'un système KKT (réduit à un complément de Schur) via le Gradient Conjugué (CG).
  • Pour $\mathcal{M}_{ub}$, la projection se simplifie en une expression en forme fermée (closed-form), éliminant le besoin de solveurs itératifs.
• Rétractions :
  • La rétraction mappe les vecteurs tangents vers la variété via une étape de miroir exponentielle cohérente avec Fisher suivie d'une projection KL.
  • Pour $\mathcal{M}_{bal}$, cela nécessite des projections de Bregman cycliques (itératives).
  • Pour $\mathcal{M}_{ub}$, la projection KL s'effondre en une mise à l'échelle de colonnes en forme fermée, éliminant ainsi les boucles internes.
• Gradient et Hessienne Riemanniens :
  • Le cadre fournit une formulation indépendante du coût où le gradient riemannien et les produits Hessienne-vecteur dépendent uniquement des dérivées euclidiennes de la fonction de coût spécifique (OT linéaire, Gromov-Wasserstein, etc.) et de la géométrie de la variété.
  • Cela permet l'utilisation de solveurs de second ordre (Trust-Region, Gradient Conjugué) qui exploitent l'information de courbure.

3. Certificat de Suffisance de Rang

L'article introduit un certificat pour la vérification de l'optimalité globale. En calculant un "gradient réduit" (matrice de slack duale) à un point critique de rang-$r$, la méthode peut vérifier si la solution est globalement optimale pour le problème de plein rang. Si le rang est insuffisant, le certificat identifie une direction pour l'augmentation du rang.

Contributions Clés

1. Formulation Riemannienne Unifiée : Formalise l'OT de bas rang équilibré et non équilibré comme des variétés lisses sous la métrique de Fisher-Rao, fournissant un cadre géométrique unique pour l'OT linéaire, le Gromov-Wasserstein (GW) et leurs variantes non équilibrées.
2. Opérations Non Équilibrées en Forme Fermée : Démontre que pour l'OT non équilibré, les projecteurs de l'espace tangent et les rétractions se réduisent à des expressions en forme fermée. Cela élimine les boucles internes de Dykstra/Bregman requises par les méthodes de pointe (ex: UB-LOT), réduisant la complexité par itération à $O(rN)$.
3. Solveurs de Second Ordre : Permet l'utilisation de solveurs sensibles à la courbure (CG, Trust-Region) qui atteignent une convergence plus rapide que la descente de miroir de premier ordre, sans nécessiter de régularisation entropique ou de réglage de la taille de pas.
4. Vérification de l'Optimalité Globale : Fournit un certificat calculable pour vérifier si une solution de bas rang est globalement optimale ou pour guider les stratégies d'augmentation de rang.

Résultats Expérimentaux

Des expériences approfondies à travers des contextes équilibrés et non équilibrés (OT linéaire, GW, Fused GW) avec des tailles de jeux de données allant jusqu'à $N=50\,000$ démontrent :

• Vitesse de Convergence : Les solveurs riemanniens convergent en un ordre de grandeur de moins d'itérations externes par rapport aux bases de descente de miroir.
• Temps d'Exécution (Wall-Clock Time) : L'élimination des boucles internes dans le cas non équilibré conduit à des accélérations significatives. Par exemple, pour l'OT non équilibré avec des valeurs aberrantes (outliers), les solveurs riemanniens convergent en quelques secondes, là où UB-LOT sature ou diverge.
• Performance : Les méthodes atteignent des valeurs d'objectif supérieures ou comparables sans le biais entropique inhérent aux bases régularisées.
• Avantage du Second Ordre : Sur des problèmes non convexes (GW, Fused GW), les méthodes Trust-Region exploitent l'information de l'Hessienne pour échapper aux minima locaux peu profonds, surpassant les méthodes de premier ordre.

Signification

L'article soutient qu'en passant d'un paradigme de descente de miroir de premier ordre à une perspective de géométrie riemannienne, les solveurs d'OT de bas rang peuvent surmonter les limites de mauvais conditionnement et de sensibilité aux hyperparamètres. La principale importance réside dans la simplification structurelle des opérations non équilibrées, qui transforme un processus itératif coûteux en une évaluation en forme fermée. Cela permet une optimisation sans régularisation et sans réglage, qui passe à l'échelle linéairement avec la taille du jeu de données tout en fournissant des garanties de convergence de second ordre rigoureuses et une vérification de l'optimalité globale. Le cadre est présenté comme une approche générale et indépendante du coût, applicable à un large éventail de problèmes de transport optimal.