Scaling Weisfeiler-Leman Expressiveness Analysis to Massive Graphs with GPUs

Sintesi Tecnica: Scalare l'Analisi dell'Espressività di Weisfeiler–Leman su Grafi Massivi con GPU

Definizione del Problema
Il test di Weisfeiler–Leman (1-WL) e la sua associata colorazione stabile sono fondamentali per comprendere il potere espressivo delle Graph Neural Networks (GNN). La colorazione stabile 1-WL fornisce un limite superiore alla capacità dei modelli GNN a passaggio di messaggi (message-passing) di distinguere tra strutture di grafi; i nodi che ricevono lo stesso colore stabile sono indistinguibili per tali reti. Tuttavia, calcolare questa colorazione stabile per grafi reali e massivi presenta due colli di bottiglia critici:

1. Esecuzione Sequenziale: Gli algoritmi classici basati sulla raffinazione delle partizioni (ad esempio, l'approccio basato su splitter con complessità $O(m \log n)$) sono intrinsecamente sequenziali e non possono sfruttare il massiccio parallelismo delle moderne architetture GPU.
2. Requisiti di Memoria Globale: Questi algoritmi richiedono che l'intero grafo risieda in memoria simultaneamente. Questa natura "globale" impedisce l'analisi di grafi su scala web (ad esempio, quelli con decine di miliardi di archi) che superano la capacità di un singolo dispositivo o anche di grandi cluster di CPU.

I tentativi esistenti di parallelizzare o linearizzare il problema spesso si affidano all'aritmetica in virgola mobile, il che introduce instabilità numerica su larga scala, o non affrontano il vincolo di memoria dei grafi più grandi della memoria del dispositivo.

Metodologia
Gli autori propongono un approccio a due fronti che scinde la scalabilità computazionale dalla scalabilità della memoria, utilizzando un'interpretazione lineare-algebrica della colorazione stabile 1-WL.

1. Raffinazione Lineare-Algebrica Randomizzata (Scalabilità Computazionale)
Invece della raffinazione sequenziale delle partizioni, gli autori riformulano il problema utilizzando la proprietà di invarianza delle partizioni equitable: una partizione è stabile se e solo se l'immagine di qualsiasi vettore costante per blocchi sotto la matrice di adiacenza rimane costante per blocchi.

• Algoritmo: Il metodo campiona iterativamente vettori casuali simmetrici sulla partizione corrente, li moltiplica per la matrice di adiacenza ($z = Aw$) e raffina la partizione in base alla simmetria del vettore risultante.
• Dominio Aritmetico: Per evitare l'instabilità numerica dell'iterazione di potenza in virgola mobile (che può causare fusioni errate su grafi grandi), l'algoritmo opera interamente all'interno dell'anello modulare $\mathbb{Z}/2^{64}\mathbb{Z}$ utilizzando l'aritmetica degli interi senza segno a 64 bit.
• Garanzia Teorica: Sfruttando il vincolo strutturale per cui le somme di riga della matrice di adiacenza sono limitate da $2^{64}/2$, gli autori derivano una garanzia di correttezza probabilistica: una sonda casuale rileva qualsiasi raffinazione mancante con una probabilità di errore di al massimo $10^{-9}$, bypassando la necessità del classico lemma di Schwartz–Zippel che richiede un campo.
• Mappatura su GPU: L'operazione principale si riduce a una singola moltiplicazione matrice-vettore sparsa (SpMV) seguita da un ordinamento lessicografico, entrambe operazioni altamente efficienti sull'hardware GPU.

2. Batching Preservante la Correttezza (Scalabilità della Memoria)
Per gestire grafi che non entrano nella memoria del dispositivo, gli autori introducono uno schema di batching che decompone l'insieme degli archi in sottoinsiemi disgiunti preservando la correttezza della colorazione stabile globale.

• Nodi Interni vs. Nodi di Confine: Per un dato batch di archi, i nodi sono classificati come "interni" (tutti gli archi uscenti rimangono all'interno del batch) o "di confine" (gli archi uscenti attraversano i confini del batch).
• Raffinazione Locale: L'algoritmo elabora ogni batch indipendentemente. Fondamentalmente, la partizione iniziale per ogni batch colloca tutti i nodi interni in un unico blocco e tutti i nodi di confine in blocchi singoli. Ciò assicura che la raffinazione unisca solo i nodi la cui informazione completa di vicinato è disponibile all'interno di quel batch.
• Contrazione Iterativa: Le raffinazioni locali vengono fuse in una partizione globale. L'algoritmo costruisce quindi un grafo quoziente (dove i nodi rappresentano i blocchi della partizione corrente) e ripete il processo.
• Terminazione: Il processo itera finché il grafo ridotto non entra in un singolo batch (ottenendo la colorazione stabile più grossolana) o finché non raggiunge un punto fisso. Se viene raggiunto un punto fisso prima che il grafo entri in memoria, il risultato è una colorazione stabile valida che è un raffinamento di quella più grossolana.

Contributi Chiave

1. Algoritmo di Raffinazione Randomizzato: Un framework Monte Carlo per la colorazione stabile 1-WL che opera in aritmetica intera con strette garanzie probabilistiche ($<10^{-9}$ di errore), evitando la perdita di precisione dei metodi in virgola mobile.
2. Schema di Batching: Una strategia di decomposizione che consente il calcolo 1-WL su grafi che eccedono la memoria del dispositivo, elaborando sottoinsiemi di archi indipendenti e fondendo i risultati tramite grafi quoziente, mantenendo provabilmente la proprietà della colorazione stabile.
3. Implementazione GPU: Un'implementazione CUDA che mappa questi algoritmi direttamente sulle primitive GPU (SpMV, sorting, reduction), ottenendo accelerazioni massicce rispetto ai baseline CPU.

Risultati Sperimentali
Gli autori hanno valutato il loro approccio su undici grafi di media scala e tre grafi su scala web (33–48 miliardi di archi).

• Correttezza: Sui grafi di media scala dove il grafo completo entra in memoria ($N_b=1$), l'algoritmo randomizzato ha recuperato l'esatta colorazione stabile più grossolana, corrispondendo ai baseline deterministici.
• Stabilità Numerica: Il metodo proposto con aritmetica intera ha distinto con successo nodi che erano stati erroneamente uniti da un precedente metodo di iterazione di potenza in virgola mobile su grafi grandi (ad esempio, il dataset it-2004 ha mostrato una discrepanza di oltre 800.000 blocchi).
• Qualità del Batching: Utilizzando più batch ($N_b=2, 4$) per simulare i vincoli di memoria, le partizioni risultanti erano empiricamente entro il 5% della colorazione stabile più grossolana. Per diversi dataset, il divario era zero.
• Prestazioni:
  • Velocità: L'implementazione su singola GPU ha ottenuto accelerazioni fino a 138× rispetto alla classica raffinazione delle partizioni CPU single-threaded e fino a 449× rispetto ai baseline CPU multicore.
  • Scalabilità: Sui grafi su scala web con oltre 30 miliardi di archi, tutti i baseline CPU (inclusa la riduzione simbolica multicore) sono andati in timeout o hanno esaurito la memoria. Il metodo proposto ha calcolato con successo le colorazioni stabili per queste istanze, ottenendo riduzioni del ~39–46% nel numero di blocchi.

Significatività e Rivendicazioni
Il lavoro afferma che questo studio è il primo a calcolare con successo colorazioni stabili 1-WL su grafi su scala web con oltre 30 miliardi di archi, un compito precedentemente impossibile a causa dei limiti di memoria e computazione. La significatività risiede nel:

• Abilitare l'Analisi dell'Espressività: Permettere l'analisi empirica dell'espressività delle GNN su dataset reali e massivi, anziché essere limitati a piccoli o medi grafi.
• Adattamento Architetturale: Dimostrare che il problema 1-WL, tradizionalmente visto come sequenziale e globale, può essere efficacemente parallelizzato e distribuito utilizzando formulazioni lineare-algebriche e strategie di batching progettate per le architetture GPU.
• Robustezza: Fornire un metodo che sia sia numericamente stabile (tramite aritmetica intera) che scalabile in termini di memoria (tramite batching), superando i specifici modi di fallimento dei metodi precedenti.

Gli autori concludono che il loro metodo consente l'analisi dell'espressività 1-WL a scale precedentemente irraggiungibili, sottolineando come la corrispondenza tra colorazione stabile e partizioni equitable offra potenziali percorsi futuri per l'applicazione di queste tecniche a compiti di machine learning su grafi e alla riduzione dei modelli per sistemi dinamici non lineari.