Scaling Weisfeiler-Leman Expressiveness Analysis to Massive Graphs with GPUs

Technische Samenvatting: Schalen van de Weisfeiler–Leman Expressiviteitsanalyse naar Massale Grafen met GPU's

Probleemstelling
De Weisfeiler–Leman (1-WL) test en de daarmee geassocieerde stabiele kleuring zijn fundamenteel voor het begrijpen van de expressieve kracht van Graph Neural Networks (GNN's). De 1-WL stabiele kleuring biedt een bovengrens aan het vermogen van message-passing GNN's om onderscheid te maken tussen grafenstructuren; knopen die dezelfde stabiele kleur ontvangen, zijn ononderscheidbaar voor dergelijke netwerken. Het berekenen van deze stabiele kleuring voor real-world, massale grafen staat echter voor twee kritieke knelpunten:

1. Sequentiële Executie: Klassieke algoritmen gebaseerd op partitieverfijning (bijv. de $O(m \log n)$ splitter-gebaseerde aanpak) zijn inherent sequentieel en kunnen de massale parallelliteit van moderne GPU-architecturen niet exploiteren.
2. Globale Geheugeneisen: Deze algoritmen vereisen dat de volledige graaf gelijktijdig in het geheugen aanwezig is. Dit "globale" karakter verhindert de analyse van web-schaal grafen (bijv. die met tientallen miljarden randen) die het capaciteitsbereik van een enkele device of zelfs grote CPU-clusters overschrijden.

Bestaande pogingen om het probleem te parallelliseren of te lineariseren vertrouwen vaak op drijvende-kommagetal rekenkunde, wat bij schaalvergroting voor numerieke instabiliteit zorgt, of adresseren het geheugengebrek van grafen die groter zijn dan het devicegeheugen niet.

Methodologie
De auteurs stellen een tweeledige aanpak voor die de rekenkundige schaalbaarheid ontkoppelt van de geheugenschaalbaarheid, door gebruik te maken van een lineair-algebraïsche interpretatie van de 1-WL stabiele kleuring.

1. Gerandomiseerde Lineair-Algebraïsche Verfijning (Rekenkundige Schaalbaarheid)
In plaats van sequentiële partitieverfijning, herformuleren de auteurs het probleem met behulp van de invariantie-eigenschap van vergelijkbare partities (equitable partitions): een partitie is stabiel als en slechts als de afbeelding van elke blok-constante vector onder de adjacentiematrix ook blok-constant blijft.

• Algoritme: De methode iteratieveven het samplen van willekeurige vectoren die symmetrisch zijn op de huidige partitie, vermenigvuldigt deze met de adjacentiematrix ($z = Aw$), en verfijnt de partitie op basis van de symmetrie van de resulterende vector.
• Rekenkundig Domein: Om de numerieke instabiliteit van floating-point machtsiteratie (die foutieve samenvoegingen op grote grafen kan veroorzaken) te vermijden, opereert het algoritme volledig binnen de modulaire ring $\mathbb{Z}/2^{64}\mathbb{Z}$ met behulp van unsigned 64-bit integer rekenkunde.
• Theoretische Garantie: Door gebruik te maken van de structurele beperking dat de rij-sommen van de adjacentiematrix begrensd zijn door $2^{64}/2$, leiden de auteurs een probabilistische correctheidsgarantie af. Een willekeurige probe detecteert elke ontbrekende verfijning met een foutkans van maximaal $10^{-9}$, waarmee de noodzaak voor de klassieke Schwartz–Zippel lemma (die een lichaam vereist) wordt omzeild.
• GPU Mapping: De kernoperatie reduceert tot een enkele sparse matrix-vector vermenigvuldiging (SpMV) gevolgd door een lexicografische sortering, beide zeer efficiënt op GPU-hardware.

2. Correctheid-Preserverende Batching (Geheugenschaalbaarheid)
Om grafen te verwerken die niet in het devicegeheugen passen, introduceren de auteurs een batching-schema dat de verzameling randen (edges) deelt in disjuncte subsets terwijl de correctheid van de globale stabiele kleuring behouden blijft.

• Innerlijke vs. Grensknopen: Voor een gegeven batch randen worden knopen geclassificeerd als "innerlijk" (alle uitgaande randen liggen binnen de batch) of "grensknoop" (uitgaande randen kruisen batchgrenzen).
• Lokale Verfijning: Het algoritme verwerkt elke batch onafhankelijk. Cruciaal is dat de initiële partitie voor elke batch alle innerlijke knopen in één blok plaatst en alle grensknopen in singleton-blokken. Dit zorgt ervoor dat verfijning alleen knopen samenvoegt waarvan de volledige buurtinformatie binnen die batch beschikbaar is.
• Iteratieve Contractie: Lokale verfijningen worden samengevoegd tot een globale partitie. Het algoritme construeert vervolgens een quotientgraaf (waarbij knopen de blokken van de huidige partitie vertegenwoordigen) en herhaalt het proces.
• Terminatie: Het proces iteratief totdat de gereduceerde graaf in een enkele batch past (wat de bewezen meest grove stabiele kleuring oplevert) of een fixpunt bereikt. Als een fixpunt wordt bereikt voordat de graaf in het geheugen past, is het resultaat een geldige stabiele kleuring die een verfijning is van de meest grove kleuring.

Belangrijkste Bijdragen

1. Gerandomiseerd Verfijningsalgoritme: Een Monte Carlo-framework voor 1-WL stabiele kleuring dat opereert in integer rekenkunde met strikte probabilistische garanties ($<10^{-9}$ fout), waardoor de precisieverlies van floating-point methoden wordt vermeden.
2. Batching-schema: Een decompositiestrategie die 1-WL berekening op grafen mogelijk maakt die het devicegeheugen overschrijden door onafhankelijke edge-subsets te verwerken en resultaten samen te voegen via quotientgrafen, waarbij de stabiele kleuringseigenschap bewezen behouden blijft.
3. GPU Implementatie: Een CUDA-gebaseerde implementatie die deze algoritmen direct mapt naar GPU-primitieven (SpMV, sortering, reductie), wat enorme versnellingen oplevert ten opzichte van CPU-baselines.

Experimentele Resultaten
De auteurs evalueerden hun aanpak op elf medium-schaal grafen en drie web-schaal grafen (33–48 miljard randen).

• Correctheid: Op medium-schaal grafen waar de volledige graaf in het geheugen past ($N_b=1$), herstelde het gerandomiseerde algoritme exact de meest grove stabiele kleuring, overeenkomend met deterministische baselines.
• Numerieke Stabiliteit: De voorgestelde integer-rekenmethode slaagde erin knopen te onderscheiden die onjuist waren samengevoegd door een eerdere floating-point machtsiteratie-methode op grote grafen (bijv. de it-2004 dataset vertoonde een discrepantie van meer dan 800.000 blokken).
• Batching Kwaliteit: Wanneer meerdere batches werden gebruikt ($N_b=2, 4$) om geheugenbeperkingen te simuleren, lag de resulterende partitie empirisch binnen 5% van de meest grove stabiele kleuring. Voor verschillende datasets was het verschil nul.
• Prestaties:
  • Snelheid: De single-GPU implementatie bereikte versnellingen van wel 138× ten opzichte van klassieke single-threaded CPU partitieverfijning en tot 449× ten opzichte van multicore CPU-baselines.
  • Schaalbaarheid: Op web-schaal grafen met meer dan 30 miljard randen, liepen alle CPU-baselines (inclusclusief multicore symbolische reductie) ofwel vast (timeout) of liepen ze zonder geheugen. De voorgestelde methode slaagde erin deze instanties te berekenen, waarbij een reductie van ~39–46% in het aantal blokken werd bereikt.

Betekenis en Claims
Het artikel claimt dat dit werk de eerste is die erin slaagt om 1-WL stabiele kleuringen te berekenen op web-schaal grafen met meer dan 30 miljard randen, een taak die voorheen onhaalbaar was door geheugen- en rekenbeperkingen. De betekenis ligt in:

• Mogelijk maken van Expressiviteitsanalyse: Het mogelijk maken van empirische analyse van GNN-expressiviteit op real-world, massale datasets in plaats van beperkt te zijn tot kleine of medium-grote grafen.
• Architecturale Fit: Het aantonen dat het 1-WL probleem, traditioneel gezien als sequentieel en globaal, effectief geparallelliseerd en gedistribueerd kan worden met behulp van lineair-algebraïsche formuleringen en batching-strategieën die op GPU-architecturen zijn afgestemd.
• Robuustheid: Het bieden van een methode die zowel numeriek stabiel is (via integer rekenkunde) als geheugenschaalbaar (via batching), waarmee de specifieke faalmodi van eerdere benaderingen worden overwonnen.

De auteurs concluderen dat hun methode 1-WL expressiviteitsanalyse mogelijk maakt op schalen die voorheen onbereikbaar waren, waarbij zij opmerken dat de correspondentie tussen stabiele kleuring en vergelijkbare partities potentieel toekomstige wegen biedt voor het toepassen van deze technieken op downstream graph ML-taken en modelreductie voor niet-lineaire dynamische systemen.