🤖 AI

FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU

FlashSinkhorn est un solveur GPU conscient des entrées-sorties pour le transport optimal entropique qui exploite la fusion et le carrelage de style FlashAttention pour réduire drastiquement le trafic mémoire HBM, atteignant des accélérations allant jusqu'à 161 fois par rapport aux bases de référence les plus avancées tout en permettant une optimisation évolutive pour des tâches de nuages de points à grande échelle.

Auteurs originaux : Felix X. -F. Ye, Xingjie Li, An Yu, Ming-Ching Chang, Linsong Chu, Davis Wertheimer

Publié 2026-05-22

📖 4 min de lecture☕ Lecture pause café

CC BY 4.0

Auteurs originaux : Felix X. -F. Ye, Xingjie Li, An Yu, Ming-Ching Chang, Linsong Chu, Davis Wertheimer

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez d'apparier deux foules immenses de personnes. Une foule se tient d'un côté d'un champ (la « source »), et l'autre de l'autre côté (la « cible »). Votre objectif est de déterminer le moyen le plus efficace d'apparier tout le monde afin que la distance totale parcourue par chacun soit minimisée. C'est un problème mathématique classique appelé Transport Optimal.

Dans l'apprentissage automatique moderne, nous ajoutons souvent une petite touche de « flou » à ce processus d'appariement pour faciliter les calculs. C'est ce qu'on appelle le Transport Optimal Entropique. Pour le résoudre, les ordinateurs utilisent une méthode appelée itérations de Sinkhorn, qui ressemble à un jeu de « patate chaude » où l'ordinateur fait passer des notes d'avant en arrière entre les deux foules, affinant les appariements encore et encore jusqu'à trouver la meilleure solution.

Le Problème : L'Embouteillage

L'article explique que, bien que cette méthode fonctionne bien pour de petites foules, elle se heurte à un mur massif lorsque les foules deviennent énormes (comme des dizaines de milliers de personnes).

Imaginez la mémoire de l'ordinateur comme une ville :

HBM (Mémoire à haute bande passante) : C'est l'autoroute principale de la ville. Elle est vaste et peut contenir beaucoup de données, mais elle est lente d'accès.
SRAM (Mémoire sur puce) : C'est un petit bureau privé ultra-rapide situé directement à l'intérieur du processeur de l'ordinateur. Il est incroyablement rapide mais très petit.

Les anciennes méthodes pour résoudre ce problème d'appariement ressemblaient à un camion de livraison qui devait faire l'aller-retour entre l'autoroute (HBM) et le bureau (SRAM) à chaque fois qu'il devait vérifier une seule paire de personnes. Comme il existe des millions de paires possibles, le camion restait coincé dans les embouteillages de l'autoroute, déplaçant constamment des données d'avant en arrière. L'ordinateur passait plus de temps à attendre les données qu'à faire les calculs.

La Solution : FlashSinkhorn

Les auteurs ont créé un nouvel outil appelé FlashSinkhorn. Ils ont réalisé que les mathématiques derrière ce problème d'appariement ressemblent exactement à celles utilisées dans les Transformers (la technologie derrière les chatbots d'IA comme celui avec qui vous parlez).

Dans les Transformers, il existe une astuce ingénieuse appelée FlashAttention qui résout un embouteillage similaire. Au lieu de faire faire l'aller-retour au camion, FlashAttention charge un « carreau » entier (un petit lot) de données dans le bureau rapide, effectue tous les calculs nécessaires là-bas, et n'écrit que le résultat final sur l'autoroute.

FlashSinkhorn adopte cette même stratégie « basée sur les carreaux » et l'applique au problème d'appariement :

Plus de Cartes Complètes : Au lieu de noter toute la carte de chaque connexion possible (qui serait trop grande pour tenir en mémoire), il calcule les connexions à la volée, un petit carreau à la fois.
La Stratégie du « Bureau » : Il garde le lot actuel de calculs dans le bureau rapide et petit (SRAM). Il met à jour les « scores d'appariement » directement là-bas, sans jamais avoir besoin de réécrire la liste intermédiaire massive sur l'autoroute lente.
Flux Continu : Il parcourt les données comme sur un tapis roulant, traitant et écartant le travail lourd au fur et à mesure, maintenant l'autoroute dégagée.

Les Résultats : Vitesse et Échelle

L'article a testé cela sur des GPU puissants (spécifiquement l'A100). Les résultats ont été spectaculaires :

Vitesse : Jusqu'à 32 fois plus rapide pour le calcul initial et jusqu'à 161 fois plus rapide pour le processus complet (y compris l'apprentissage des erreurs) par rapport aux meilleures méthodes en ligne existantes.
Mémoire : Alors que les anciennes méthodes plantaient (manque de mémoire) en essayant d'apparier des foules de 30 000 personnes, FlashSinkhorn pouvait gérer 50 000 personnes facilement car il n'essayait jamais de stocker toute la carte d'un coup.
Utilisation Réelle : Ils ont montré que cela fonctionne sur des tâches réelles comme la comparaison d'énormes ensembles de données (comme des milliers d'images) et la résolution de problèmes de régression complexes où l'ordre des données est mélangé.

La Conclusion

FlashSinkhorn est comme passer d'un camion de livraison coincé dans les embouteillages à un drone haute vitesse. Il ne change pas la destination (la réponse mathématique reste exacte), mais il change la façon dont les données sont déplacées. En gardant le travail lourd à l'intérieur du « bureau » rapide de l'ordinateur et en n'utilisant l'autoroute lente que pour les résultats finaux, il rend la résolution de problèmes d'appariement massifs pratique et rapide, transformant une tâche qui prenait autrefois des heures ou faisait planter l'ordinateur en quelque chose qui ne prend que quelques secondes.

Résumé technique : FlashSinkhorn : Transport Optimal Entropique conscient des E/S sur GPU

Énoncé du problème

Le Transport Optimal Entropique (EOT), résolu via des itérations de Sinkhorn, est un outil fondamental en apprentissage automatique moderne pour comparer des mesures de probabilité, allant de la modélisation générative à l'alignement de distributions. Cependant, la mise à l'échelle de l'EOT vers de grands nuages de points (par exemple, $n \approx 70\,000$ ) sur GPU reste un goulot d'étranglement persistant.

Les solveurs GPU existants font face à deux inefficacités principales :

Implémentations tensorisées : Elles matérialisent la matrice de coût dense $n \times m$ , entraînant un trafic mémoire HBM (High-Bandwidth Memory) quadratique ($O(nm)$ lectures/écritures) et des problèmes de capacité mémoire provoquant des erreurs de dépassement de mémoire (OOM) à mesure que $n$ augmente.
Backends en ligne : Bien qu'ils évitent de stocker la matrice dense en évaluant les interactions à la volée (par exemple via KeOps ou OTT-JAX), ils reposent sur des noyaux génériques de réduction par tuiles (map-reduce). Ces noyaux souffrent d'une fusion limitée, empêchant l'exploitation des structures GEMM des cœurs tensoriels et résultant en une utilisation du calcul sous-optimale et un temps d'exécution élevé malgré une réduction de l'utilisation de la mémoire.

Le goulot d'étranglement système central est que les méthodes actuelles sont soit limitées par la mémoire (en raison d'un trafic HBM excessif), soit limitées par le calcul avec une efficacité d'instruction médiocre, échouant à s'adapter aux résolutions répétées à grande échelle requises dans les pipelines en aval.

Méthodologie

Les auteurs introduisent FlashSinkhorn, un solveur GPU conscient des E/S pour l'EOT avec un coût euclidien au carré. La méthodologie établit un parallèle direct entre les mises à jour stabilisées de Sinkhorn dans le domaine logarithmique et le mécanisme d'attention des Transformers, exploitant les principes de FlashAttention pour optimiser les motifs d'accès mémoire.

1. Reformulation sous forme d'attention

L'insight théorique clé (Proposition 1) est que, pour un coût euclidien au carré, les mises à jour stabilisées de Sinkhorn dans le domaine logarithmique peuvent être réécrites comme des réductions LogSumExp (LSE) ligne par ligne d'une matrice de scores de produit scalaire biaisée.

Plus précisément, la mise à jour $f_i \leftarrow -\varepsilon \text{LSE}_j [(g_j - C_{ij})/\varepsilon + \log b_j]$ est équivalente au calcul du LSE d'une matrice $S_X = (QK^\top + \mathbf{1}(\hat{g} + \delta))/\varepsilon$ , où $Q=\sqrt{2}X$ , $K=\sqrt{2}Y$ , et $\delta, \gamma$ sont des biais de poids log précalculés.
Cette reformulation transforme l'itération de Sinkhorn en la même statistique de normalisation utilisée dans l'attention par produit scalaire mis à l'échelle, permettant l'application de techniques de streaming conscientes des E/S.

2. Noyaux de streaming et tuilage

FlashSinkhorn implémente cette reformulation en utilisant des noyaux Triton fusionnés qui emploient un modèle mémoire à deux niveaux (HBM lent et SRAM rapide sur puce) :

Pas de matérialisation : La matrice de coût $n \times m$ n'est jamais matérialisée dans la HBM. À la place, des tuiles de $Q$ et $K$ sont streamées.
LSE en ligne : Les noyaux maintiennent des statistiques de somme-exp et de maximum ligne par ligne en cours d'exécution sur puce (SRAM) au fur et à mesure qu'ils traitent les tuiles de la matrice de scores. Cela permet le calcul de réductions LSE exactes avec un nombre significativement inférieur d'accès HBM.
Fusion : Toute la mise à jour (calcul du score, ajout de biais, réduction LSE et mise à jour potentielle) est fusionnée en un seul passage de noyau. Cela élimine les écritures intermédiaires vers la HBM et réduit la surcharge de lancement de noyaux.
Parcours stationnaire par ligne : Contrairement à FlashAttention, qui boucle souvent sur les clés/valeurs dans la boucle externe, FlashSinkhorn adopte un parcours stationnaire par ligne (fixant un bloc de ligne de requête dans la SRAM et streamant tous les blocs de clés) pour mieux correspondre à la nature de réduction ligne par ligne des mises à jour de Sinkhorn.

3. Différentiation et optimisation d'ordre supérieur

Le papier étend cette stratégie consciente des E/S à la différentiation et aux méthodes d'ordre supérieur :

Application de la matrice de transport : La projection barycentrique et le calcul du gradient sont exprimés comme des opérations de type attention (moyennes pondérées), permettant leur calcul via les mêmes noyaux de streaming sans matérialiser la matrice de transport $P^*$ .
Produits Hessian-Vecteur (HVP) : Les auteurs dérivent un oracle HVP par streaming qui décompose le produit Hessian-vecteur en une séquence de produits transport-vecteur et transport-matrice. Cela évite de matérialiser le tenseur Hessian $O(n^2 d^2)$ ou la matrice de transport $O(nm)$, réduisant la mémoire de travail à $O((n+m)d)$ .

Contributions clés

Mise à jour de Sinkhorn sous forme d'attention : Les auteurs expriment les mises à jour de Sinkhorn stabilisées comme des réductions LSE de produit scalaire biaisé, permettant un streaming exact de style FlashAttention pour l'EOT.
Solveur EOT conscient des E/S : Ils fournissent des noyaux Triton fusionnés pour les itérations de Sinkhorn en streaming et des opérateurs de streaming sans matrice de transport pour la différentiation, réduisant substantiellement le volume d'E/S HBM.
Optimisation d'ordre supérieur évolutive : Ils introduisent des HVP par streaming qui permettent des méthodes Newton-CG et de région de confiance évolutives pour les pertes basées sur l'EOT, auparavant irréalisables à grande échelle en raison des contraintes mémoire.
Évaluation à grande échelle : Le papier démontre des accélérations substantielles et une efficacité mémoire sur des GPU A100 à travers des benchmarks synthétiques et des tâches en aval telles que la Distance de Jeu de Données par Transport Optimal (OTDD) et la régression mélangée.

Résultats expérimentaux

Évalué sur des GPU NVIDIA A100-80 Go, FlashSinkhorn démontre des gains de performance significatifs par rapport aux références de l'état de l'art (GeomLoss avec backends KeOps/Tensorisées et OTT-JAX) :

Vitesse :
- Passage avant : Atteint jusqu'à 32× d'accélération par rapport à KeOps et 5,1× par rapport à OTT-JAX.
- Extrémité à extrémité (Avant + Arrière) : Atteint jusqu'à 161× d'accélération par rapport à KeOps à haute dimension ( $d=512$ ).
- HVP : Atteint 3–25× d'accélération selon la dimensionnalité, rendant Newton-CG pratique à des échelles où d'autres méthodes dépassent les délais.
Efficacité mémoire :
- FlashSinkhorn maintient une complexité mémoire linéaire $O((n+m)d)$ , tandis que les références tensorisées évoluent en $O(nm)$ ou $O(n^2)$ .
- Il traite avec succès des nuages de points avec $n=50\,000$ et des dimensions allant jusqu'à $d=1024$ , là où les méthodes tensorisées manquent de mémoire et où les références en ligne (KeOps) dépassent les limites de temps.
Tâches en aval :
- OTDD : FlashSinkhorn met à l'échelle le calcul de la distance OTDD et le flot de gradient jusqu'à $n=60\,000$ (MNIST $\leftrightarrow$ Fashion-MNIST) avec une mémoire de pointe inférieure à 1 Go, tandis que les références tensorisées échouent par manque de mémoire (OOM) pour $n > 20\,000$ .
- Évasion de selle : Dans la régression linéaire mélangée, le HVP par streaming permet de surveiller la plus petite valeur propre du Hessien pour détecter les points de selle, permettant à un optimiseur hybride Adam-Newton de s'échapper des selles et de converger rapidement.

Importance et revendications

Le papier revendique que FlashSinkhorn rend l'EOT à grande échelle pratique dans le régime limité par la mémoire. En réécrivant les mises à jour de Sinkhorn pour correspondre à la structure consciente des E/S de FlashAttention, les auteurs comblent le fossé entre le transport optimal et les techniques modernes d'accélération des transformers.

L'importance réside dans :

Optimisation au niveau système : Les gains sont attribués à une spécialisation au niveau noyau (fusion et tuilage) plutôt qu'à des changements algorithmiques de l'itération de Sinkhorn elle-même.
Permettre de nouvelles charges de travail : La réduction de l'empreinte mémoire et du temps d'exécution permet l'utilisation de l'EOT dans des tâches en aval (comme la distance de jeu de données et la différentiation implicite) à des échelles précédemment considérées comme prohibitives sur le plan computationnel.
Reproductibilité : Les auteurs publient une implémentation open-source pour faciliter la recherche future sur le transport optimal évolutif.

Les auteurs notent modestement que cette approche s'applique actuellement aux coûts ayant une structure spécifique (euclidien au carré et distance cosinus après normalisation). Les coûts sans cette structure (par exemple, euclidien brut avec racines carrées ou coûts neuronaux appris) sont identifiés comme un travail futur.