Fixed-Point Neural Optimal Transport without Implicit Differentiation

Résumé technique : Transport optimal neuronal à point fixe sans différenciation implicite

Énoncé du problème
Le transport optimal (OT) cherche une application entre des distributions de probabilité qui minimise un coût de transport prescrit. Bien que des approches récentes d'apprentissage profond aient permis de développer des solveurs OT évolutifs, elles reposent fréquemment sur une optimisation min–max adversaire (par exemple, les WGAN) ou sur des architectures multi-réseaux impliquant des applications de transport et des potentiels duaux. Ces méthodes souffrent souvent d'instabilité lors de l'entraînement, de sensibilité aux hyperparamètres et de difficultés de convergence, en particulier dans des contextes de grande dimension. Les approches alternatives basées sur des flots continus nécessitent la résolution d'équations différentielles ordinaires (EDO), entraînant une surcharge computationnelle significative, tandis que les formulations de Hamilton–Jacobi–Bellman (HJB) peinent à être évolutives en grande dimension. De plus, de nombreuses méthodes existantes requièrent des réseaux auxiliaires ou des procédures d'optimisation alternée pour assurer la faisabilité duale.

Méthodologie
Les auteurs proposent un cadre à réseau unique pour le transport optimal quadratique qui élimine l'entraînement adversaire et la différenciation implicite. Le cœur de la méthode est une formulation neuronale implicite dérivée du problème dual de Kantorovich.

1. Paramétrisation unique du potentiel : Le potentiel dual $\phi$ est paramétré comme l'opposé d'une fonction convexe $g_\theta$ représentée par un seul réseau neuronal ($\phi_\theta(x) = -g_\theta(x)$). Cela exploite le théorème de Brenier, qui stipule que l'application de transport optimale est le gradient d'un potentiel convexe.
2. Reformulation par point fixe proximal : Sous cette paramétrisation, la $c$-transform (un composant clé de l'objectif dual) se réduit à un problème d'optimisation proximal fortement convexe :
   $$(-g_\theta)^c(z) = \inf_{y \in \mathbb{R}^d} \left\{ \frac{1}{2}\|y - z\|^2 + g_\theta(y) \right\}$$
   Le minimiseur $y^*_\theta(z)$ de ce problème définit l'application de transport inverse ($T_{\nu \to \mu}(z) = y^*_\theta(z)$), tandis que l'application directe est donnée par $T_{\mu \to \nu}(x) = x + \nabla g_\theta(x)$.
3. Entraînement par itération de point fixe : L'objectif d'entraînement consiste à minimiser la perte duale négative. Crucialement, la minimisation interne (l'opérateur proximal) est résolue à l'aide d'itérations de point fixe standard (par exemple, des étapes de descente de gradient) plutôt que par une optimisation déroulée ou une différenciation implicite.
4. Calcul du gradient sans différenciation implicite : Une insight théorique clé est que le gradient de la perte par rapport aux paramètres du réseau $\theta$ ne nécessite pas de rétropropagation à travers les itérations de point fixe. En raison des conditions d'optimalité du premier ordre de l'opérateur proximal, les termes impliquant la dérivée du point fixe par rapport à $\theta$ ($dy^*/d\theta$) s'annulent exactement. Par conséquent, le gradient ne dépend que des dérivées explicites de $g_\theta$ évaluées aux échantillons d'entrée et aux solutions de point fixe approximatives. Cela permet un entraînement efficace utilisant la rétropropagation standard.
5. Convergence : Les auteurs établissent que la Descente de Gradient Stochastique (SGD) converge vers un point stationnaire même lorsque le problème interne de point fixe est résolu de manière approximative (dans une tolérance spécifiée), à condition que l'erreur d'approximation soit bornée.
6. Extension conditionnelle par classe : Le cadre s'étend naturellement au OT conditionnel par classe en conditionnant le potentiel $g_\theta(x, k)$ sur les étiquettes de classe $k$, permettant à un seul réseau de gérer plusieurs tâches de transport simultanément.

Contributions clés

• Reformulation à réseau unique : Les auteurs dérivent une formulation du OT quadratique qui ne nécessite qu'un seul réseau neuronal pour retrouver les applications de transport directe et inverse, éliminant le besoin de configurations min–max adversaires ou de réseaux conjugués auxiliaires.
• Entraînement sans différenciation implicite : Ils démontrent que l'objectif peut être différencié sans différenciation implicite, en s'appuyant sur l'annulation exacte des termes induite par les conditions d'optimalité de l'opérateur proximal.
• Garanties de convergence : L'article fournit des preuves théoriques de la convergence de la SGD vers un point stationnaire sous des évaluations de point fixe inexactes, répondant à la réalité pratique selon laquelle la boucle interne n'est pas résolue avec une précision machine.
• Cadre unifié : La méthode récupère simultanément les applications directe et inverse et s'étend naturellement aux contextes conditionnels par classe sans modifications architecturales.

Résultats expérimentaux
La méthode proposée a été évaluée sur des benchmarks gaussiens de grande dimension, des jeux de données physiques réels (UCI) et des tâches de traduction d'images (MNIST vers Fashion-MNIST).

• Gaussiennes de grande dimension : Sur des benchmarks gaussiens avec des dimensions allant jusqu'à $d=64$, la méthode a atteint un pourcentage de variance inexpliquée (UVP) nettement inférieur par rapport à des bases de référence telles que NOT, MM-v1, MM:R, NCF et WGAN-QC. Elle a maintenu une précision stable à mesure que la dimension augmentait, tandis que les bases de référence montraient une dégradation.
• Efficacité computationnelle : La méthode a démontré une efficacité computationnelle et mémoire favorable, nécessitant moins de paramètres et moins de temps d'entraînement que les méthodes basées sur des flots continus (FFJORD, RNODE) et d'autres approches neuronales de OT.
• Données réelles : Sur divers jeux de données UCI (Power, Gas, HEPMASS, MINIBOONE) et des patches d'images BSDS300, la méthode a atteint des performances de correspondance de distribution compétitives ou supérieures (mesurées par la Divergence Maximale Moyenne) avec des tailles de modèles substantiellement plus petites.
• Traduction d'images : Dans les tâches de traduction d'images conditionnelles par classe, la méthode a obtenu la plus faible distance Fréchet Inception (FID) et la plus haute précision de classification parmi les bases de référence, indiquant une préservation supérieure de la structure de classe et du contenu sémantique.
• Études d'ablation : Les expériences ont montré que la méthode est robuste à la tolérance d'arrêt du solveur de point fixe ; une inexactitude modérée dans la boucle interne offre un compromis favorable entre précision et vitesse. De plus, bien que l'imposition d'une convexité stricte via une projection des poids ait stabilisé les itérations de point fixe, elle a dégradé la précision globale du transport, suggérant qu'une paramétrisation non convexe (sans projection stricte) offre un meilleur équilibre entre expressivité et stabilité.

Portée et affirmations
L'article affirme que ce cadre offre une alternative stable, efficace et évolutive aux méthodes neuronales de OT existantes. En éliminant l'entraînement adversaire et la surcharge computationnelle des solveurs d'EDO ou de la résolution d'équations HJB sur tout le domaine, la méthode fournit un mécanisme « sans causalité » pour imposer les conditions d'optimalité. Les auteurs affirment que leur approche est largement applicable, capable de gérer des problèmes de grande dimension et conditionnels par classe tout en récupérant des applications de transport précises dans les deux sens. Ce travail se positionne comme une étape vers la simplification du paysage d'optimisation du OT, s'éloignant des formulations min–max complexes pour se diriger vers un problème de minimisation à objectif unique, théoriquement fondé et pratiquement efficace.