Technische Samenvatting: cuRegOT

Probleemstelling

Optimal Transport (OT) is een fundamenteel hulpmiddel in het moderne machine learning voor het vergelijken van waarschijnlijkheidsmaten, en wordt veel gebruikt in domeinadaptatie, generatieve modellering en computer vision. Het oplossen van het discrete OT-probleem op grote schaal is echter computationeel onhaalbaar vanwege de superkubische complexiteit van standaard lineaire programmeringsoplossers. Hoewel entropische regularisatie (Cuturi, 2013) het probleem transformeert in een glad, onbeperkt dual optimalisatieprobleem dat oplosbaar is via het Sinkhorn-algoritme, lijdt dit laatste vaak aan trage convergentie bij uitdagende problemen.

Recente vooruitgang heeft tweede-orde en quasi-Newton-methoden geïntroduceerd (zoals SNS, SSNS en het Sparse-Plus-Low-Rank of SPLR-algoritme) die superieure convergentiesnelheden bieden door de structuur van de Hessiaan-matrix te benutten. Deze methoden staan echter voor een kritieke bottleneck wanneer ze worden ingezet op GPU's: de noodzaak om grote, verspreide lineaire systemen op te lossen. Specifiek zijn de symbolische analyse en herordening stappen die nodig zijn voor een verspreide Cholesky-decompositie inherent sequentieel, moeilijk te paralleliseren op SIMT-architecturen, en worden ze doorgaans uitgevoerd op de CPU. Dit creëert een synchronisatiebottleneck waarbij de GPU inactief blijft, wat de efficiëntie van anderszins krachtige quasi-Newton-oplossers ernstig belemmert.

Methodologie

Het artikel presenteert cuRegOT, een hoogpresterende GPU-oplosser op maat gemaakt voor entropisch-geregulariseerde OT, gebaseerd op het SPLR-algoritme. Om de synchronisatiebottleneck tussen CPU en GPU en de beperkingen van geheugenbandbreedte te overwinnen, stellen de auteurs drie specifieke algoritmische en architecturale optimalisaties voor:

1. Geamortiseerde Symbolische Analyse

Standaard SPLR evalueert het verspreidheidspatroon van de Hessiaan-benadering ($H_\Omega$) bij elke iteratie opnieuw, wat dure CPU-gebaseerde symbolische analyse triggert. De auteurs merken op dat het verspreidheidspatroon (bepaald door de top-$k$ waarden van het transportplan $T$) zich langzaam ontwikkelt over iteraties heen.

• Strategie: De oplosser voert de CPU-gebaseerde symbolische analyse en matrixherordening slechts één keer elke $S$ iteraties uit (bijvoorbeeld $S=10$). Voor de tussenliggende iteraties wordt het vastgestelde verspreidheidspatroon hergebruikt, waarbij alleen de numerieke waarden van de niet-nul-elementen worden bijgewerkt.
• Theoretische Garantie: De auteurs bewijzen dat het hergebruiken van het verspreidheidspatroon het eigenwaardebereik van de Hessiaan-benadering niet uitbreidt. Bijgevolg blijft het voorwaartse getal begrensd, waardoor de convergentie-eigenschappen van de quasi-Newton-methode behouden blijven.

2. Collaboratief CPU-GPU Computing (Asynchrone Kandidaatgeneratie)

Zelfs met amortisatie blijft de GPU inactief tijdens de periodieke CPU-gebaseerde symbolische analyse-fases.

• Strategie: De auteurs introduceren een asynchroon mechanisme waarbij, terwijl de CPU symbolische analyse uitvoert, de GPU gelijktijdig lichtgewicht Sinkhorn-iteraties uitvoert op de huidige dual variabelen. Dit genereert een "kandidaat"-iteraat ($x^s_k$).
• Selectie: Zodra de CPU de analyse afrondt en de GPU de primaire quasi-Newton-stap voltooit (waarbij $x^q_k$ wordt gegenereerd), evalueert de oplosser de objectieve functie en de normen van de gradiënt voor beide kandidaten. Het iteraat dat de grootste verbetering oplevert, wordt geselecteerd.
• Voordeel: Deze aanpak maskeert CPU-latentie en biedt extra verkenning van de oplossingsruimte zonder de wandkloktijd te verhogen. Theoretische analyse bevestigt dat het invoegen van deze kandidaatselectiestappen de globale convergentie van het oorspronkelijke algoritme behoudt.

3. Gefuseerde CUDA-Kernel voor Gradiëntevaluatie

Naïeve gradiëntberekening omvat meerdere doorgangen over de kostenmatrix $M$ en het transportplan $T$, wat leidt tot ernstige knelpunten in de geheugenbandbreedte op GPU's.

• Strategie: Een gefuseerde CUDA-kernel is ontworpen om elementgewijze berekening, rijreductie en kolomreductie uit te voeren in een enkele doorgang over de data.
• Implementatie: De kernel maakt gebruik van een grid-stride-lus en maximaliseert datahergebruik in registers en gedeeld geheugen. Het maakt gebruik van warp-niveau shuffle-intrinsieke functies voor rijreducties (nul gebruik van gedeeld geheugen) en buffers in gedeeld geheugen voor kolomreducties, waardoor globale geheugenschrijfsels en atomische concurrentie worden geminimaliseerd.

Belangrijkste Bijdragen

1. cuRegOT-bibliotheek: Een hoogpresterende GPU-oplosser voor entropisch-geregulariseerde OT die de SPLR quasi-Newton-methode succesvol aanpast aan GPU-hardware.
2. Algoritmische Innovaties:
   • Een geamortiseerde symbolische analyse-strategie die CPU-overhead vermindert door verspreidheidspatronen te hergebruiken.
   • Een asynchroon kandidaatgeneratie-mechanisme dat de inactieve tijd van de GPU tijdens CPU-analyse benut om convergentie te versnellen.
   • Een gefuseerde gradiëntkernel die de overhead van geheugen-I/O aanzienlijk vermindert.
3. Theoretische Garanties: Rigoureuze bewijzen die aantonen dat de voorgestelde wijzigingen (patroonhergebruik en kandidaatselectie) de globale convergentie en de lineaire convergentiesnelheid van het oorspronkelijke SPLR-algoritme behouden.

Experimentele Resultaten

De auteurs hebben cuRegOT geëvalueerd tegenover state-of-the-art GPU-gebaseerde oplossers, waaronder POT (Sinkhorn, Greenkhorn), OTT-JAX (Sinkhorn, Anderson-acceleratie) en AccSinkhorn, op synthetische data en het CIFAR-10-dataset.

• Prestaties: cuRegOT bereikt consistent aanzienlijk snellere convergentie dan alle baselines, met name wanneer hoge precisie vereist is.
• Schaalbaarheid: Het prestatievoordeel van cuRegOT neemt toe met de probleemgrootte (getest tot $6400 \times 4800$ en $5000 \times 5000$), wat de schaalbaarheid voor grootschalige toepassingen valideert.
• Ablatiestudie: Experimenten waarbij de geamortiseerde analyse ("-A") of de kandidaatgeneratie ("-S") afzonderlijk werden verwijderd, toonden langzamere convergentie, wat bevestigt dat beide componenten essentieel zijn voor de prestaties van de oplosser.
• Baselines: Hoewel Anderson-acceleratie in sommige gevallen snelheidswinst bood voor standaard Sinkhorn, bleef het kampen met trage convergentie naar hoge precisie in vergelijking met cuRegOT. Greenkhorn bleek te traag voor een praktische vergelijking.

Betekenis

Het artikel adresseert een kritieke kloof in de implementatie van geavanceerde optimalisatiemethoden voor Optimal Transport op moderne hardware. Door de efficiëntiekloof te overbruggen tussen de theoretische convergentiesnelheid van quasi-Newton-methoden en de praktische beperkingen van GPU-uitvoering (specifiek met betrekking tot verspreide lineaire oplossers), maakt cuRegOT het mogelijk om grootschalige, hoogprecisie OT-problemen efficiënt op te lossen. Het werk demonstreert dat met zorgvuldig algoritmisch ontwerp – specifiek het amortiseren van seriële CPU-taken en het fuseren van geheugen-gebonden operaties – tweede-orde methoden haalbaar kunnen worden gemaakt en superieur kunnen zijn aan eerste-orde baselines op GPU-versnellers.