Résumé Technique : cuRegOT

Énoncé du Problème

Le Transport Optimal (OT) est un outil fondamental en apprentissage automatique moderne pour la comparaison de mesures de probabilité, largement utilisé dans l'adaptation de domaine, la modélisation générative et la vision par ordinateur. Cependant, la résolution du problème d'OT discret à grande échelle est prohibitivement coûteuse en calcul en raison de la complexité super-cubique des solveurs de programmation linéaire standard. Bien que la régularisation entropique (Cuturi, 2013) transforme le problème en un problème d'optimisation duale lisse et sans contrainte, soluble via l'algorithme de Sinkhorn, ce dernier souffre souvent d'une convergence lente sur des problèmes difficiles.

Des avancées récentes ont introduit des méthodes d'ordre deux et quasi-Newton (par exemple, SNS, SSNS, et l'algorithme Sparse-Plus-Low-Rank ou SPLR) qui offrent des taux de convergence supérieurs en exploitant la structure de la matrice hessienne. Cependant, ces méthodes rencontrent un goulot d'étranglement critique lors du déploiement sur GPU : la nécessité de résoudre de grands systèmes linéaires creux. Plus précisément, les étapes d'analyse symbolique et de réordonnancement nécessaires à la décomposition de Cholesky creuse sont intrinsèquement séquentielles, difficiles à paralléliser sur des architectures SIMT, et généralement exécutées sur le CPU. Cela crée un goulot d'étranglement de synchronisation où le GPU reste inactif, entravant sévèrement l'efficacité de solveurs quasi-Newton par ailleurs puissants.

Méthodologie

L'article présente cuRegOT, un solveur haute performance sur GPU conçu pour l'OT régularisé entropiquement, basé sur l'algorithme SPLR. Pour surmonter le goulot d'étranglement de synchronisation CPU-GPU et les limitations de bande passante mémoire, les auteurs proposent trois optimisations spécifiques, tant algorithmiques qu'architecturales :

1. Analyse Symbolique Amortie

Le SPLR standard réévalue le motif de creusité de l'approximation hessienne ($H_\Omega$) à chaque itération, déclenchant une analyse symbolique coûteuse sur CPU. Les auteurs observent que le motif de creusité (déterminé par les $k$ premières entrées du plan de transport $T$) évolue lentement au fil des itérations.

• Stratégie : Le solveur effectue l'analyse symbolique sur CPU et le réordonnancement de la matrice une seule fois tous les $S$ itérations (par exemple, $S=10$). Pour les itérations intermédiaires, le motif de creusité établi est réutilisé, en mettant à jour uniquement les valeurs numériques des éléments non nuls.
• Garantie Théorique : Les auteurs prouvent que la réutilisation du motif de creusité n'élargit pas la plage de valeurs propres de l'approximation hessienne. Par conséquent, le nombre de conditionnement reste borné, garantissant que les propriétés de convergence de la méthode quasi-Newton sont préservées.

2. Calcul Collaboratif CPU-GPU (Génération Asynchrone de Candidats)

Même avec l'amortissement, le GPU reste inactif pendant les phases périodiques d'analyse symbolique sur CPU.

• Stratégie : Les auteurs introduisent un mécanisme asynchrone où, pendant que le CPU effectue l'analyse symbolique, le GPU exécute simultanément des itérations légères de Sinkhorn sur les variables duales courantes. Cela génère une itérée « candidate » ($x^s_k$).
• Sélection : Une fois le CPU terminé l'analyse et le GPU achevé l'étape quasi-Newton principale (générant $x^q_k$), le solveur évalue la fonction objectif et les normes du gradient pour les deux candidats. L'itérée produisant la plus grande amélioration est sélectionnée.
• Bénéfice : Cette approche masque la latence du CPU et offre une exploration supplémentaire de l'espace des solutions sans augmenter le temps d'exécution réel. L'analyse théorique confirme que l'insertion de ces étapes de sélection de candidats préserve la convergence globale de l'algorithme original.

3. noyau CUDA Fusionné pour l'Évaluation du Gradient

Le calcul naïf du gradient implique plusieurs passages sur la matrice de coût $M$ et le plan de transport $T$, entraînant des goulots d'étranglement sévères de bande passante mémoire sur les GPU.

• Stratégie : Un noyau CUDA fusionné est conçu pour effectuer le calcul élément par élément, la réduction de ligne et la réduction de colonne en un seul passage sur les données.
• Implémentation : Le noyau utilise une boucle à pas de grille (grid-stride loop) et maximise la réutilisation des données dans les registres et la mémoire partagée. Il emploie des intrinsèques de mélange au niveau du warp (warp-level shuffle) pour les réductions de ligne (sans utilisation de mémoire partagée) et des tampons de mémoire partagée pour les réductions de colonne, minimisant ainsi les écritures en mémoire globale et la contention atomique.

Contributions Clés

1. Bibliothèque cuRegOT : Un solveur haute performance sur GPU pour l'OT régularisé entropiquement qui adapte avec succès la méthode quasi-Newton SPLR au matériel GPU.
2. Innovations Algorithmiques :
   • Une stratégie d'analyse symbolique amortie qui réduit la surcharge CPU en réutilisant les motifs de creusité.
   • Un mécanisme de génération asynchrone de candidats qui utilise le temps d'inactivité du GPU pendant l'analyse CPU pour accélérer la convergence.
   • Un noyau de gradient fusionné qui réduit considérablement la surcharge d'E/S mémoire.
3. Garanties Théoriques : Des preuves rigoureuses démontrant que les modifications proposées (réutilisation de motif et sélection de candidats) maintiennent la convergence globale et le taux de convergence linéaire de l'algorithme SPLR original.

Résultats Expérimentaux

Les auteurs ont évalué cuRegOT par rapport à des solveurs basés sur GPU de l'état de l'art, notamment POT (Sinkhorn, Greenkhorn), OTT-JAX (Sinkhorn, accélération d'Anderson) et AccSinkhorn, sur des données synthétiques et le jeu de données CIFAR-10.

• Performance : cuRegOT atteint systématiquement une convergence nettement plus rapide que toutes les lignes de base, en particulier lorsque une haute précision est requise.
• Évolutivité : L'avantage de performance de cuRegOT augmente avec la taille du problème (testé jusqu'à $6400 \times 4800$ et $5000 \times 5000$), validant son évolutivité pour des applications à grande échelle.
• Étude d'Ablation : Des expériences retirant l'analyse amortie ("-A") ou la génération de candidats ("-S") individuellement ont montré une convergence plus lente, confirmant que les deux composants sont essentiels à la performance du solveur.
• Lignes de Base : Bien que l'accélération d'Anderson ait fourni des accélérations pour Sinkhorn standard dans certains cas, elle a toujours fait face à une convergence lente vers une haute précision par rapport à cuRegOT. Greenkhorn s'est révélé trop lent pour une comparaison pratique.

Importance

L'article comble une lacune critique dans le déploiement de méthodes d'optimisation avancées pour le Transport Optimal sur du matériel moderne. En comblant l'écart d'efficacité entre la vitesse de convergence théorique des méthodes quasi-Newton et les limitations pratiques de l'exécution sur GPU (spécifiquement concernant les solveurs linéaires creux), cuRegOT permet la résolution efficace de problèmes OT à grande échelle et haute précision. Ce travail démontre qu'avec une conception algorithmique soignée — spécifiquement l'amortissement des tâches CPU sérielles et la fusion des opérations liées à la mémoire — les méthodes d'ordre deux peuvent devenir viables et supérieures aux lignes de base d'ordre premier sur les accélérateurs GPU.