Sintesi Tecnica: GraphGP: Processi Gaussiani Scalabili con l'Approssimazione di Vecchia

Definizione del Problema

I processi gaussiani (GP) forniscono un quadro rigoroso per la modellazione spaziale e la quantificazione dell'incertezza, ma affrontano gravi limitazioni di scalabilità. Le implementazioni naive comportano costi computazionali di $O(N^3)$ e requisiti di memoria di $O(N^2)$, rendendoli impraticabili per dataset con milioni o miliardi di punti. Sebbene i metodi esistenti offrano soluzioni parziali, essi presentano significativi compromessi:

• Metodi basati su punti inducenti: riducono la complessità a $O(NM^2)$, ma faticano quando il numero di variabili inducenti $M$ deve essere elevato.
• Trasformate di Fourier Veloci (FFT): raggiungono una scalabilità $O(N \log N)$ su griglie regolari, ma sono difficili da applicare a distribuzioni di punti arbitrarie con ampi intervalli dinamici.
• L'approssimazione di Vecchia: offre una scalabilità $O(Nk^3)$ condizionando ogni punto solo su $k \ll N$ predecessori, ma l'esecuzione efficiente di questo processo per grandi $N$ su distribuzioni di punti arbitrarie è stata una sfida complessa.
• Iterative Charted Refinement (ICR): scala a centinaia di milioni di parametri su una singola GPU, ma restringe i punti a una griglia gerarchica, limitandone la flessibilità.

La sfida principale affrontata è come determinare rapidamente le relazioni di condizionamento appropriate e parallelizzare il calcolo risultante per distribuzioni di punti arbitrarie senza sacrificare la scalabilità ottenuta con i metodi basati su griglia.

Metodologia

GraphGP introduce un algoritmo GPU che combina la flessibilità dell'approssimazione di Vecchia con un'elevata parallelizzazione delle prestazioni. La metodologia principale si basa su due innovazioni fondamentali: una nuova strategia di ordinamento dei punti e un'implementazione CUDA personalizzata.

1. Approssimazione di Vecchia

L'algoritmo fattorizza la densità congiunta $p(v)$ di $N$ punti come un prodotto di probabilità condizionate:
$$p(v) \approx \prod_{i=1}^{N} p(v_i | v_{c(i)})$$
dove $c(i)$ è un insieme di $k$ punti di condizionamento precedenti. Inveve di utilizzare tutti i punti precedenti (il che sarebbe esatto ma costoso), GraphGP condiziona ogni punto solo sui suoi $k$ vicini più prossimi tra i punti precedenti. Ciò riduce le matrici di condizionamento a una dimensione $k \times k$, consentendo una complessità temporale di $O(Nk^3)$ e una complessità di memoria di $O(Nk^2)$.

2. Ordinamento $k$-d Tree con Bit-Reversal

L'efficienza dell'approssimazione di Vecchia dipende fortemente dall'ordine di generazione dei punti. Gli ordini standard dei $k$-d tree raggruppano punti spazialmente vicini, creando lunghe catene di dipendenza che ostacolano la parallelizzazione. GraphGP propone un ordinamento bit-reversed:

• Struttura ad Albero: I punti sono organizzati in un $k$-d tree bilanciato a sinistra.
• Permutazione Bit-Reversed: L'ordine di generazione viene modificato in modo che tutti i figli sinistri siano inseriti prima di qualsiasi figlio destro a ogni livello. Questo separa efficacemente i punti spazialmente vicini nella sequenza di generazione.
• Benefici: Questo ordinamento assicura che i punti consecutivi nella sequenza di generazione siano spazialmente distanti, rendendoli probabilmente indipendenti condizionatamente. Di conseguenza, grandi batch di punti possono essere generati in parallelo.
• Costruzione dei Batch: L'algoritmo analizza il grafo aciclico diretto (DAG) implicito nelle relazioni di vicinato per calcolare la "profondità" (il percorso più lungo dalla radice) di ogni nodo. I punti allo stesso livello di profondità sono condizionatamente indipendenti e vengono elaborati in parallelo. Per garantire l'accuratezza ed evitare piccoli batch iniziali, i primi $n_0$ punti vengono generati utilizzando una fattorizzazione di Cholesky densa.

3. Implementazione CUDA e Ottimizzazioni

Sebbene esista un'implementazione baseline in JAX, GraphGP utilizza un'implementazione CUDA personalizzata per superare i colli di bottiglia di memoria e prestazioni:

• Calcolo basato su Registri: Poiché le matrici di condizionamento $k \times k$ sono piccole ($k \lesssim 64$) e a breve durata, esse vengono calcolate on-the-fly nei registri della GPU piuttosto che essere allocate nella memoria globale. Ciò evita i limiti di banda della memoria.
• Operazioni In-Place: L'implementazione utilizza fattorizzazioni matriciali e risoluzioni in-place, memorizzando le matrici in formato triangolare inferiore per minimizzare l'uso dei registri.
• Differenziabilità: Per supportare l'inferenza bayesiana e l'adattamento dei modelli, gli autori hanno implementato derivate forward e backward in CUDA utilizzando le derivate della decomposizione di Cholesky. Queste sono registrate come primitive JAX personalizzate, permettendo a GraphGP di funzionare senza problemi all'interno di modelli differenziabili più ampi.
• Costruzione del Grafo in Parallelo: La ricerca dei vicini utilizza algoritmi GPU basati su sorting parallelo adattati per l'ordine bit-reversed personalizzato. Gli aggiornamenti della profondità vengono eseguiti iterativamente su tutti i punti in parallelo anziché sequenzialmente, il che è più veloce per i grafi poco profondi tipici di questa applicazione.

Contributi Chiave

1. Ordinamento $k$-d Tree con Bit-Reversal: Un nuovo ordine di generazione che massimizza il parallelismo dei batch mantenendo capacità efficienti di ricerca dei vicini all'interno di una struttura ad albero.
2. Implementazione CUDA Differenziabile: Un'implementazione ad alte prestazioni e a basso consumo di memoria, sostanzialmente più veloce di una baseline pura in JAX, dotata di differenziazione automatica per le derivate dei parametri del kernel.
3. Scalabilità a quasi un miliardo di parametri: La combinazione della strategia di ordinamento e delle ottimizzazioni hardware permette all'algoritmo di scalare a quasi $10^9$ parametri su una singola GPU con requisiti di tempo e memoria lineari.
4. Componenti completi per l'inferenza: Il sistema fornisce generazione forward, applicazione inversa, calcolo del log-determinante e derivate dei parametri del kernel, consentendo una completa inferenza bayesiana.

Risultati

• Prestazioni: I benchmark su una GPU H100 dimostrano che l'implementazione CUDA è più di un ordine di grandezza più veloce della baseline JAX per il passaggio forward.
• Scalabilità: Il passaggio forward raggiunge una scalabilità lineare rispetto al numero di punti $N$.
• Efficienza della Memoria: Calcolando le matrici nei registri, l'impronta di memoria è significativamente ridotta rispetto alle strategie di allocazione in memoria globale.
• Accuratezza: Le valutazioni tramite divergenza di Kullback–Leibler mostrano che, sebbene l'ordine di generazione influenzi l'accuratezza, l'effetto è minore rispetto al numero di vicini $k$ utilizzati. Il proposto ordinamento ottiene i batch paralleli più grandi con una perdita di accuratezza dell'approssimazione trascurabile rispetto agli ordini standard di alberi o casuali.
• Confronto: A differenza di GpGpU, che implementa solo le operazioni inverse e di log-determinante e manca della ricerca dei vicini su GPU, GraphGP supporta l'intero passaggio forward e distribuzioni di punti arbitrarie. A differenza di ICR, non richiede che i punti siano limitati a una griglia.

Significato

GraphGP abilita l'applicazione pratica dei processi gaussiani alla mappatura tridimensionale del mezzo interstellare e di altri campi della fisica che richiedono l'inferenza su campi con risoluzione disomogenea. Rimuovendo i vincoli dei metodi basati su griglia e i colli di bottiglia della memoria dei GP naivi, consente la generazione di mappe ad alta risoluzione con requisiti di memoria inferiori. Questo lavoro colma il divario tra la flessibilità dell'approssimazione di Vecchia su punti arbitrari e la massiccia scalabilità precedentemente ottenibile solo su griglie strutturate, aprendo la strada a una nuova generazione di strumenti di inferenza bayesiana per dati spaziali su larga scala.