Technische Samenvatting: FlashSinkhorn: IO-bewuste Entropische Optimale Transport op GPU

Probleemstelling

Entropische Optimale Transport (EOT), opgelost via Sinkhorn-iteraties, is een fundamenteel hulpmiddel in modern machine learning voor het vergelijken van waarschijnlijkheidsmaten, variërend van generatieve modellering tot distributie-uitlijning. Het schalen van EOT naar grote puntwolken (bijv. $n \approx 70.000$) op GPU's blijft echter een aanhoudende knelpunt.

Bestaande GPU-oplossers kampen met twee primaire inefficiënties:

1. Getensoriseerde Implementaties: Deze materialiseren de dichte $n \times m$ kostenmatrix, wat leidt tot kwadratisch High-Bandwidth Memory (HBM)-verkeer ($O(nm)$ leest/schrijft) en geheugencapaciteitsproblemen die Out-of-Memory (OOM)-fouten veroorzaken naarmate $n$ groeit.
2. Online Backends: Hoewel deze het opslaan van de dichte matrix vermijden door interacties on-the-fly te evalueren (bijv. via KeOps of OTT-JAX), vertrouwen ze op generieke getegelde map-reduce-reductiekernen. Deze kernen lijden aan beperkte fusie, wat het benutten van tensor-core GEMM-structuren verhindert en resulteert in suboptimale rekenuitbating en hoge runtime ondanks verminderd geheugengebruik.

De kern van het systeemknelpunt is dat huidige methoden óf geheugenbeperkt zijn (door excessief HBM-verkeer) óf rekengedwongen met slechte instructie-efficiëntie, waardoor ze niet schalen naar de herhaalde grootschalige oplossingen die nodig zijn in downstream-pijplijnen.

Methodologie

De auteurs introduceren FlashSinkhorn, een IO-bewuste GPU-oplosser voor EOT met kwadratische Euclidische kosten. De methodologie trekt een directe parallel tussen gestabiliseerde log-domein Sinkhorn-updates en het attentiemechanisme in Transformers, waarbij de principes van FlashAttention worden benut om geheugentoegangspatronen te optimaliseren.

1. Herformulering in de vorm van Attentie

Het belangrijkste theoretische inzicht (Propositie 1) is dat voor kwadratische Euclidische kosten, de gestabiliseerde log-domein Sinkhorn-updates herschreven kunnen worden als rijgewijze LogSumExp (LSE)-reducties van een bevooroordeelde dot-product scorematrix.

• Specifiek is de update $f_i \leftarrow -\varepsilon \text{LSE}_j [(g_j - C_{ij})/\varepsilon + \log b_j]$ equivalent aan het berekenen van de LSE van een matrix $S_X = (QK^\top + \mathbf{1}(\hat{g} + \delta))/\varepsilon$, waarbij $Q=\sqrt{2}X$, $K=\sqrt{2}Y$, en $\delta, \gamma$ vooraf berekende log-gewichtbevooroordeelingen zijn.
• Deze herformulering transformeert de Sinkhorn-iteratie naar exact dezelfde normalisatiestatistiek die wordt gebruikt in geschaalde dot-product attentie, waardoor de toepassing van IO-bewuste streamingtechnieken mogelijk wordt.

2. Streamingkernen en Tegeling

FlashSinkhorn implementeert deze herformulering met behulp van gefuseerde Triton-kernen die een tweeledig geheugenmodel gebruiken (traag HBM en snel on-chip SRAM):

• Geen Materialisatie: De $n \times m$ kostenmatrix wordt nooit gematerialiseerd in HBM. In plaats daarvan worden tegels van $Q$ en $K$ gestreamd.
• Online LSE: De kernen houden lopende rijgewijze max- en sum-exp-statistieken bij on-chip (SRAM) terwijl ze tegels van de scorematrix verwerken. Dit maakt het berekenen van exacte LSE-reducties mogelijk met aanzienlijk minder HBM-toegang.
• Fusie: De volledige update (scoreberekening, bevooroordeeling toevoegen, LSE-reductie en potentiële update) wordt gefuseerd tot een enkele kernpassage. Dit elimineert tussentijdse schrijfsels naar HBM en verlaagt de overhead van kernstarten.
• Rij-stationaire Traversie: In tegenstelling tot FlashAttention, dat vaak over sleutels/waarden in de buitenste lus loopt, adopteert FlashSinkhorn een rij-stationaire traversie (een query-rijblok fixeren in SRAM en alle sleutelblokken streamen) om beter aan te sluiten bij de rijgewijze reductie-aard van Sinkhorn-updates.

3. Differentiatie en Optimalisatie van de Tweede Orde

Het artikel breidt deze IO-bewuste strategie uit tot differentiatie en methoden van de tweede orde:

• Toepassing van Transportmatrix: De barycentrische projectie en gradientberekening worden uitgedrukt als attentie-achtige operaties (gewogen gemiddelden), waardoor ze via dezelfde streamingkernen kunnen worden berekend zonder de transportmatrix $P^*$ te materialiseren.
• Hessiaan-Vector Producten (HVP): De auteurs leiden een streaming HVP-orakel af dat het Hessiaan-vector product decomposeert in een reeks transport-vector en transport-matrix producten. Dit voorkomt het materialiseren van de $O(n^2 d^2)$ Hessiaanstensor of de $O(nm)$ transportmatrix, waardoor het werkgeheugen wordt gereduceerd tot $O((n+m)d)$.

Belangrijkste Bijdragen

1. Sinkhorn-update in de vorm van Attentie: De auteurs drukken gestabiliseerde Sinkhorn-updates uit als bevooroordeelde dot-product LSE-reducties, waardoor exacte FlashAttention-stijl streaming voor EOT mogelijk wordt.
2. IO-bewuste EOT-oplosser: Zij leveren gefuseerde Triton-kernen voor gestreamde Sinkhorn-iteraties en transportmatrix-vrije streamingoperatoren voor differentiatie, wat het HBM-IO-volume aanzienlijk reduceert.
3. Schalbare Optimalisatie van de Tweede Orde: Zij introduceren streaming HVP's die schaalbare Newton-CG en trust-region methoden mogelijk maken voor EOT-gebaseerde verliezen, die eerder onhaalbaar waren op grote schaal vanwege geheugenbeperkingen.
4. Grootschalige Evaluatie: Het artikel toont aanzienlijke snelheidswinsten en geheugenefficiëntie aan op A100 GPU's, zowel op synthetische benchmarks als op downstream-taken zoals Optimal Transport Dataset Distance (OTDD) en geschudde regressie.

Experimentele Resultaten

Geëvalueerd op NVIDIA A100-80GB GPU's, demonstreert FlashSinkhorn significante prestatiewinsten ten opzichte van state-of-the-art baselines (GeomLoss met KeOps/getensoriseerde backends en OTT-JAX):

• Snelheid:
  • Forward Pass: Bereikt tot 32× snelheidswinst ten opzichte van KeOps en 5,1× ten opzichte van OTT-JAX.
  • End-to-End (Forward + Backward): Bereikt tot 161× snelheidswinst ten opzichte van KeOps bij hoge dimensies ($d=512$).
  • HVP: Bereikt 3–25× snelheidswinst afhankelijk van de dimensionaliteit, waardoor Newton-CG praktisch wordt op schalen waar andere methoden time-outen.
• Geheugenefficiëntie:
  • FlashSinkhorn behoudt lineaire geheugencomplexiteit $O((n+m)d)$, terwijl getensoriseerde baselines schalen als $O(nm)$ of $O(n^2)$.
  • Het verwerkt succesvol puntwolken met $n=50.000$ en dimensies tot $d=1024$, waarbij getensoriseerde methoden het geheugen opraken en online baselines (KeOps) de tijdslimieten overschrijden.
• Downstream-taken:
  • OTDD: FlashSinkhorn schaalt OTDD-afstandsberekening en stroom van gradiënten naar $n=60.000$ (MNIST $\leftrightarrow$ Fashion-MNIST) met een piekgeheugen onder 1 GB, terwijl getensoriseerde baselines OOM krijgen bij $n > 20.000$.
  • Zadelontsnapping: In geschudde lineaire regressie maakt de streaming HVP het mogelijk de minimale Hessiaan-eigenwaarde te monitoren om zadelpunten te detecteren, waardoor een hybride Adam-Newton-optimizer zadelpunten kan ontvluchten en snel kan convergeren.

Betekenis en Claims

Het artikel claimt dat FlashSinkhorn grootschalige EOT praktisch maakt in het geheugenbeperkte regime. Door Sinkhorn-updates te herschrijven om te matchen met de IO-bewuste structuur van FlashAttention, overbruggen de auteurs de kloof tussen optimale transport en moderne transformer-versnellingstechnieken.

De betekenis ligt in:

• Systeemniveau-optimalisatie: De winsten worden toegeschreven aan kernniveau-specialisatie (fusie en tegeling) in plaats van algoritmische wijzigingen aan de Sinkhorn-iteratie zelf.
• Mogelijk maken van nieuwe werklasten: De reductie in geheugenvoetafdruk en runtime maakt het gebruik van EOT in downstream-taken (zoals datasetafstand en impliciete differentiatie) mogelijk op schalen die eerder als rekenkundig prohibitief werden beschouwd.
• Reproduceerbaarheid: De auteurs geven een open-source implementatie vrij om verder onderzoek in schaalbaar optimale transport te faciliteren.

De auteurs merken bescheiden op dat deze aanpak momenteel van toepassing is op kosten met een specifieke structuur (kwadratische Euclidische en cosinusafstand na normalisatie). Kosten zonder deze structuur (bijv. ruwe Euclidische met vierkantswortels of geleerde neurale kosten) worden geïdentificeerd als toekomstig werk.