Technische Samenvatting: GraphGP: Schaalbare Gaussische Processen met de Vecchia-benadering

Probleemstelling

Gaussische processen (GP's) bieden een principieel kader voor ruimtelijke modellering en onzekerheidskwantificatie, maar kampen met ernstige schaalbaarheidsbeperkingen. Naïeve implementaties brengen $O(N^3)$ computationele kosten en $O(N^2)$ geheugeneisen met zich mee, wat ze onpraktisch maakt voor datasets met miljoenen of miljarden punten. Hoewel bestaande methoden gedeeltelijke oplossingen bieden, hebben zij significante afruilpunten:

• Inducing-point methoden reduceren de complexiteit naar $O(NM^2)$, maar worstelen wanneer het aantal inducerende variabelen $M$ groot moet zijn.
• Fast Fourier Transforms (FFT's) bereiken $O(N \log N)$ schaling op regelmatige roosters, maar zijn moeilijk toe te passen op willekeurige puntverdelingen met grote dynamische bereiken.
• Vecchia's benadering biedt $O(Nk^3)$ schaling door elk punt te conditioneren op slechts $k \ll N$ voorafgaande punten, maar het efficiënt uitvoeren hiervan voor grote $N$ op willekeurige puntverdelingen is een uitdaging gebleken.
• Iterative Charted Refinement (ICR) schaalt naar honderden miljoenen parameters op een enkele GPU, maar beperkt punten tot een hiërarchisch rooster, wat de flexibiliteit beperkt.

De primaire uitdaging die wordt aangepakt, is hoe men passende conditioneringsrelaties snel kan bepalen en de resulterende berekening kan parallelliseren voor willekeurige puntverdelingen zonder de schaalbaarheid op te offeren die door roostergebaseerde methoden wordt bereikt.

Methodologie

GraphGP introduceert een GPU-algoritme dat de flexibiliteit van de Vecchia-benadering combineert met hoogwaardige parallellisatie. De kernmethodologie rust op twee belangrijke innovaties: een nieuwe ordeningsstrategie voor punten en een aangepaste CUDA-implementatie.

1. Vecchia-benadering

Het algoritme factoriseert de gezamenlijke dichtheid $p(v)$ van $N$ punten als een product van conditionele waarschijnlijkheden:
$$p(v) \approx \prod_{i=1}^{N} p(v_i | v_{c(i)})$$
waarbij $c(i)$ een verzameling van $k$ voorafgaande conditioneringspunten is. In plaats van alle voorgaande punten te gebruiken (wat exact maar duur zou zijn), conditioneert GraphGP elk punt alleen op zijn $k$ dichtstbijzijnde buren onder de voorafgaande punten. Dit reduceert de conditioneringsmatrices tot de grootte $k \times k$, wat een $O(Nk^3)$ tijdcomplexiteit en $O(Nk^2)$ geheugencomplexiteit mogelijk maakt.

2. Bit-Reversed $k$-d Tree Ordening

De efficiëntie van de Vecchia-benadering hangt sterk af van de generatievolgorde van de punten. Standaard $k$-d boom-ordes groeperen ruimtelijk nabijgelegen punten, wat lange afhankelijkheidsketens creëert die parallellisatie hinderen. GraphGP stelt een bit-reversed ordering voor:

• Boomstructuur: Punten worden georganiseerd in een links-gebalanceerde $k$-d boom.
• Bit-Reversal Permutatie: De generatievolgorde wordt gewijzigd zodat alle linker kinderen worden ingevoegd voordat een rechter kind wordt ingevoegd op elk niveau. Dit spreidt ruimtelijk nabijgelegen punten effectief uit in de generatiesequentie.
• Voordelen: Deze ordening zorgt ervoor dat opeenvolgende punten in de generatiesequentie ruimtelijk ver uit elkaar liggen, waardoor ze waarschijnlijk conditioneel onafhankelijk zijn. Bijgevolg kunnen grote batches punten parallel worden gegenereerd.
• Batch Constructie: Het algoritme analyseert de gerichte acyclische graaf (DAG) die wordt geïmpliceerd door de buurrelaties om de "diepte" (langste pad vanaf de wortel) van elke node te berekenen. Punten op dezelfde diepte zijn conditioneel onafhankelijk en worden parallel verwerkt. Om nauwkeurigheid te garanderen en kleine initiële batches te voorkomen, worden de eerste $n_0$ punten gegenereerd met behulp van een dense Cholesky-factorisatie.

3. CUDA-implementatie en Optimalisaties

Hoewel er een baseline-implementatie in JAX bestaat, gebruikt GraphGP een aangepaste CUDA-implementatie om geheugen- en prestatieknelpunten te overwinnen:

• Register-gebaseerde Berekening: Omdat de $k \times k$ conditioneringsmatrices klein en kortstondig zijn ($k \lesssim 64$), worden ze on-the-fly berekend in GPU-registers in plaats van in het globale geheugen te worden gealloceerd. Dit voorkomt beperkingen in de geheugenbandbreedte.
• In-place Operaties: De implementatie maakt gebruik van in-place matrixfactorisaties en oplossingen, waarbij matrices in een lager-drihoeksformaat worden opgeslagen om het registergebruik te minimaliseren.
• Differentieerbaarheid: Om Bayesiaanse inferentie en modelpassing te ondersteunen, hebben de auteurs voorwaartse en achterwaartse afgeleiden geïmplementeerd in CUDA met behulp van Cholesky-decompositie-afgeleiden. Deze zijn geregistreerd als aangepaste JAX-primitieven, waardoor GraphGP naadloos kan functioneren binnen grotere differentieerbare modellen.
• Parallelle Graafconstructie: Zoekopdrachten naar buren maken gebruik van parallelle, op sorteren gebaseerde GPU-algoritmen die zijn aangepast voor de aangepaste bit-reversed orde. Diepte-updates worden iteratief over alle punten in parallel uitgevoerd in plaats van sequentieel, wat sneller is voor de ondiepe grafen die typerend zijn voor deze toepassing.

Belangrijkste Bijdragen

1. Bit-Reversed $k$-d Tree Ordering: Een nieuwe generatievolgorde die de batch-parallellisatie maximaliseert terwijl de efficiënte zoekcapaciteiten voor buren binnen een boomstructuur behouden blijven.
2. Differentieerbare CUDA-implementatie: Een hoogwaardige, geheugenefficiënte implementatie die aanzienlijk sneller is dan een pure JAX-baseline, voorzien van automatische differentiatie voor de afgeleiden van kernelparameters.
3. Schaalbaarheid naar Bijna een Miljard Parameters: De combinatie van de ordeningsstrategie en hardware-optimalisaties stelt het algoritme in staat om bijna $10^9$ parameters te schalen op een enkele GPU met lineaire tijd- en geheugenvereisten.
4. Volledige Inferentie Bouwstenen: Het systeem biedt voorwaartse generatie, inverse toepassing, log-determinant berekening en kernelparameter-afgeleiden, wat volledige Bayesiaanse inferentie mogelijk maakt.

Resultaten

• Prestaties: Benchmarks op een H100 GPU laten zien dat de CUDA-implementatie meer dan een orde van grootte sneller is dan de JAX-baseline voor de forward pass.
• Schaling: De forward pass bereikt lineaire schaling met betrekking tot het aantal punten $N$.
• Geheugenefficiëntie: Door matrices in registers te berekenen, wordt de geheugenvoetafdruk aanzienlijk verminderd vergeleken met strategieën voor globale geheugenallocatie.
• Nauwkeurigheid: Evaluaties met behulp van de Kullback–Leibler divergentie tonen aan dat hoewel de generatievolgorde de nauwkeurigheid beïnvloedt, het effect gering is vergeleken met het aantal buren $k$ dat wordt gebruikt. De voorgestelde ordening bereikt de grootste parallelle batches met een verwaarloosbaar verlies in benaderingsnauwkeurigheid vergeleken met standaard boom- of willekeurige ordeningen.
• Vergelijking: In tegen tegenstelling tot GpGpU, dat alleen inverse en log-determinant operaties implementeert en geen GPU-nabuurschapzoekopdrachten heeft, ondersteunt GraphGP de volledige forward pass en willekeurige puntverdelingen. In tegenstelling tot ICR vereist het niet dat punten beperkt zijn tot een rooster.

Betekenis

GraphGP maakt de praktische toepassing van Gaussische processen mogelijk voor de driedimensionale mapping van het interstellaire medium en andere natuurkundige wetenschappelijke velden die inferentie op velden met inhomogene resolutie vereisen. Door de beperkingen van roostergebaseerde methoden en de geheugenknelpunten van naïeve GP's weg te nemen, maakt het de generatie van hoogresolutiekaarten mogelijk met lagere geheugenvereisten. Dit werk overbrugt de kloof tussen de flexibiliteit van de Vecchia-benadering op willekeurige punten en de enorme schaalbaarheid die voorheen alleen bereikbaar was op gestructureerde roosters, wat de weg vrijmaakt voor een nieuwe generatie Bayesiaanse inferentietools voor grootschalige ruimtelijke data.