Scaling Weisfeiler-Leman Expressiveness Analysis to Massive Graphs with GPUs

Technical Summary: Scaling Weisfeiler–Leman Expressiveness Analysis to Massive Graphs with GPUs

Problem Statement
The Weisfeiler–Leman (1-WL) test and its associated stable coloring are fundamental to understanding the expressive power of Graph Neural Networks (GNNs). The 1-WL stable coloring provides an upper bound on the ability of message-passing GNNs to distinguish between graph structures; nodes receiving the same stable color are indistinguishable to such networks. However, computing this stable coloring for real-world, massive graphs faces two critical bottlenecks:

1. Sequential Execution: Classical algorithms based on partition refinement (e.g., the $O(m \log n)$ splitter-based approach) are inherently sequential and cannot exploit the massive parallelism of modern GPU architectures.
2. Global Memory Requirements: These algorithms require the entire graph to reside in memory simultaneously. This "global" nature prevents the analysis of web-scale graphs (e.g., those with tens of billions of edges) that exceed the capacity of single devices or even large CPU clusters.

Existing attempts to parallelize or linearize the problem often rely on floating-point arithmetic, which introduces numerical instability at scale, or fail to address the memory constraint of graphs larger than device memory.

Methodology
The authors propose a two-pronged approach that decouples compute scalability from memory scalability, utilizing a linear-algebraic interpretation of the 1-WL stable coloring.

1. Randomized Linear-Algebraic Refinement (Compute Scalability)
Instead of sequential partition refinement, the authors reformulate the problem using the invariance property of equitable partitions: a partition is stable if and only if the image of any block-constant vector under the adjacency matrix remains block-constant.

• Algorithm: The method iteratively samples random vectors symmetric on the current partition, multiplies them by the adjacency matrix ($z = Aw$), and refines the partition based on the symmetry of the resulting vector.
• Arithmetic Domain: To avoid the numerical instability of floating-point power iteration (which can cause incorrect merges on large graphs), the algorithm operates entirely within the modular ring $\mathbb{Z}/2^{64}\mathbb{Z}$ using unsigned 64-bit integer arithmetic.
• Theoretical Guarantee: By leveraging the structural constraint that row sums of the adjacency matrix are bounded by $2^{64}/2$, the authors derive a probabilistic correctness guarantee. A random probe detects any missing refinement with a failure probability of at most $10^{-9}$, bypassing the need for the classical Schwartz–Zippel lemma which requires a field.
• GPU Mapping: The core operation reduces to a single sparse matrix-vector multiplication (SpMV) followed by a lexicographic sort, both of which are highly efficient on GPU hardware.

2. Correctness-Preserving Batching (Memory Scalability)
To handle graphs that do not fit in device memory, the authors introduce a batching scheme that decomposes the edge set into disjoint subsets while preserving the correctness of the global stable coloring.

• Inner vs. Boundary Nodes: For a given batch of edges, nodes are classified as "inner" (all outgoing edges lie within the batch) or "boundary" (outgoing edges cross batch boundaries).
• Local Refinement: The algorithm processes each batch independently. Crucially, the initial partition for each batch places all inner nodes in a single block and all boundary nodes in singleton blocks. This ensures that refinement only merges nodes whose complete neighborhood information is available within that batch.
• Iterative Contraction: Local refinements are merged into a global partition. The algorithm then constructs a quotient graph (where nodes represent blocks of the current partition) and repeats the process.
• Termination: The process iterates until the reduced graph fits into a single batch (yielding the provably coarsest stable coloring) or reaches a fixed point. If a fixed point is reached before the graph fits in memory, the result is a valid stable coloring that is a refinement of the coarsest one.

Key Contributions

1. Randomized Refinement Algorithm: A Monte Carlo framework for 1-WL stable coloring that operates in integer arithmetic with tight probabilistic guarantees ($<10^{-9}$ error), avoiding the precision loss of floating-point methods.
2. Batching Scheme: A decomposition strategy that allows 1-WL computation on graphs exceeding device memory by processing independent edge subsets and merging results via quotient graphs, provably maintaining the stable coloring property.
3. GPU Implementation: A CUDA-based implementation that maps these algorithms directly to GPU primitives (SpMV, sorting, reduction), achieving massive speedups over CPU baselines.

Experimental Results
The authors evaluated their approach on eleven medium-scale graphs and three web-scale graphs (33–48 billion edges).

• Correctness: On medium-scale graphs where the full graph fits in memory ($N_b=1$), the randomized algorithm recovered the exact coarsest stable coloring, matching deterministic baselines.
• Numerical Stability: The proposed integer-arithmetic method successfully distinguished nodes that were incorrectly merged by a prior floating-point power-iteration method on large graphs (e.g., the it-2004 dataset showed a discrepancy of over 800,000 blocks).
• Batching Quality: When using multiple batches ($N_b=2, 4$) to simulate memory constraints, the resulting partitions were empirically within 5% of the coarsest stable coloring. For several datasets, the gap was zero.
• Performance:
  • Speed: The single-GPU implementation achieved speedups of up to 138× over classical single-threaded CPU partition refinement and up to 449× over multicore CPU baselines.
  • Scalability: On web-scale graphs with over 30 billion edges, all CPU baselines (including multicore symbolic reduction) either timed out or ran out of memory. The proposed method successfully computed stable colorings for these instances, achieving reductions of ~39–46% in the number of blocks.

Significance and Claims
The paper claims that this work is the first to successfully compute 1-WL stable colorings on web-scale graphs with over 30 billion edges, a task previously infeasible due to memory and compute limitations. The significance lies in:

• Enabling Expressiveness Analysis: Allowing empirical analysis of GNN expressiveness on real-world, massive datasets rather than being limited to small or medium graphs.
• Architectural Fit: Demonstrating that the 1-WL problem, traditionally viewed as sequential and global, can be effectively parallelized and distributed using linear-algebraic formulations and batching strategies tailored for GPU architectures.
• Robustness: Providing a method that is both numerically stable (via integer arithmetic) and memory-scalable (via batching), overcoming the specific failure modes of prior approaches.

The authors conclude that their method enables 1-WL expressiveness analysis at scales previously unattainable, while noting that the correspondence between stable coloring and equitable partitions offers potential future avenues for applying these techniques to downstream graph ML tasks and model reduction for nonlinear dynamical systems.