Scaling Weisfeiler-Leman Expressiveness Analysis to Massive Graphs with GPUs

Résumé Technique : Mise à l'échelle de l'analyse de l'expressivité de Weisfeiler–Leman vers les graphes massifs avec des GPU

Énoncé du Problème
Le test de Weisfeiler–Leman (1-WL) et son coloration stable associée sont fondamentaux pour comprendre la puissance expressive des réseaux de neurones sur graphes (GNN). La coloration stable 1-WL fournit une borne supérieure sur la capacité des GNN à transmettre des messages à distinguer les structures de graphes ; les nœuds recevant la même couleur stable sont indiscernables pour de tels réseaux. Cependant, le calcul de cette coloration stable pour des graphes réels et massifs se heurte à deux goulots d'étranglement critiques :

1. Exécution Séquentielle : Les algorithmes classiques basés sur le raffinement de partitions (par exemple, l'approche par splitters en $O(m \log n)$) sont intrinsèquement séquentiels et ne peuvent pas exploiter le parallélisme massif des architectures GPU modernes.
2. Exigences de Mémoire Globale : Ces algorithmes nécessitent que l'intégralité du graphe réside en mémoire simultanément. Cette nature « globale » empêche l'analyse de graphes à l'échelle du Web (par exemple, ceux possédant des dizaines de milliards d'arêtes) qui dépassent la capacité de dispositifs uniques ou même de clusters de CPU de grande taille.

Les tentatives existantes pour paralléliser ou linéariser le problème reposent souvent sur l'arithmétique en virgule flottante, ce qui introduit une instabilité numérique à grande échelle, ou échouent à traiter la contrainte de mémoire des graphes plus grands que la mémoire du dispositif.

Méthodologie
Les auteurs proposent une approche à deux volets qui découple la scalabilité de calcul de la scalabilité de mémoire, en utilisant une interprétation algébrique linéaire de la coloration stable 1-WL.

1. Raffinement Algébrique Linéaire Randomisé (Scalabilité de Calcul)
Au lieu du raffinement de partition séquentiel, les auteurs reformulent le problème en utilisant la propriété d'invariance des partitions équitables : une partition est stable si et seulement si l'image de tout vecteur constant sur un bloc sous la matrice d'adjacence reste constante sur ce bloc.

• Algorithme : La méthode échantillonne de manière itérative des vecteurs aléatoires symétriques sur la partition actuelle, les multiplie par la matrice d'adjacence ($z = Aw$), et affine la partition en fonction de la symétrie du vecteur résultant.
• Domaine Arithmétique : Pour éviter l'instabilité numérique de l'itération de puissance en virgule flottante (qui peut provoquer des fusions incorrectes sur de grands graphes), l'algorithme opère entièrement dans l'anneau modulaire $\mathbb{Z}/2^{64}\mathbb{Z}$ en utilisant l'arithmétique des entiers non signés de 64 bits.
• Garantie Théorique : En exploitant la contrainte structurelle que les sommes de lignes de la matrice d'adjacence sont bornées par $2^{64}/2$, les auteurs dérivent une garantie de correction probabiliste. Une sonde aléatoire détecte toute absence de raffinement avec une probabilité d'échec d'au plus $10^{-9}$, contournant ainsi le besoin du lemme de Schwartz–Zippel classique qui nécessite un corps.
• Mapping GPU : L'opération centrale se réduit à une seule multiplication matrice-vecteur creuse (SpMV) suivie d'un tri lexicographique, lesquels sont hautement efficaces sur le matériel GPU.

2. Batching Préservant la Correction (Scalabilité de Mémoire)
Pour gérer les graphes qui ne tiennent pas dans la mémoire du dispositif, les auteurs introduisent un schéma de batching (traitement par lots) qui décompose l'ensemble des arêtes en sous-ensembles disjoints tout en préservant la correction de la coloration stable globale.

• Nœuds Internes vs Nœuds de Bordure : Pour un lot d'arêtes donné, les nœuds sont classés comme « internes » (toutes les arêtes sortantes se trouvent dans le lot) ou « de bordure » (les arêtes sortantes traversent les limites du lot).
• Raffinement Local : L'algorithme traite chaque lot indépendamment. Crucialement, la partition initiale pour chaque lot place tous les nœuds internes dans un seul bloc et tous les nœuds de bordure dans des blocs singuliers. Cela garantit que le raffinement ne fusionne que les nœuds dont l'information complète de voisinage est disponible au sein de ce lot.
• Contraction Itérative : Les raffinements locaux sont fusionnés en une partition globale. L'algorithme construit ensuite un graphe quotient (où les nœuds représentent les blocs de la partition actuelle) et répète le processus.
• Terminaison : Le processus itère jusqu'à ce que le graphe réduit tienne dans un seul lot (obtenant ainsi la coloration stable la plus grossière) ou atteigne un point fixe. Si un point fixe est atteint avant que le graphe ne tienne en mémoire, le résultat est une coloration stable valide qui est un raffinement de la plus grossière.

Contributions Clés

1. Algorithme de Raffinement Randomisé : Un cadre de Monte Carlo pour la coloration stable 1-WL qui opère en arithmétique entière avec des garanties probabilistes serrées ($<10^{-9}$ d'erreur), évitant la perte de précision des méthodes d'itération de puissance en virgule flottante.
2. Schéma de Batching : Une stratégie de décomposition permettant le calcul 1-WL sur des graphes dépassant la mémoire du dispositif en traitant des sous-ensembles d'arêtes indépendants et en fusionnant les résultats via des graphes quotients, préservant de manière prouvée la propriété de coloration stable.
3. Implémentation GPU : Une implémentation CUDA qui mappe ces algorithmes directement sur les primitives GPU (SpMV, tri, réduction), atteignant des accélérations massives par rapport aux bases de référence CPU.

Résultats Expérimentaux
Les auteurs ont évalué leur approche sur onze graphes de taille moyenne et trois graphes de taille Web (33 à 48 milliards d'arêtes).

• Correction : Sur les graphes de taille moyenne où le graphe complet tient en mémoire ($N_b=1$), l'algorithme randomisé a récupéré la coloration stable exacte, correspondant aux bases de référence déterministes.
• Stabilité Numérique : La méthode proposée en arithmétique entière a réussi à distinguer des nœuds qui étaient incorrectement fusionnés par une méthode antérieure d'itération de puissance en virgule flottante sur de grands graphes (par exemple, le jeu de données it-2004 a montré une divergence de plus de 800 000 blocs).
• Qualité du Batching : Lors de l'utilisation de plusieurs lots ($N_b=2, 4$) pour simuler les contraintes de mémoire, les partitions résultantes étaient empiriquement proches à 5 % de la coloration stable la plus grossière. Pour plusieurs jeux de données, l'écart était nul.
• Performance :
  • Vitesse : L'implémentation sur un seul GPU a atteint des accélérations allant jusqu'à 138× par rapport au raffinement de partition classique sur un seul thread CPU et jusqu'à 449× par rapport aux bases de référence multicœurs CPU.
  • Scalabilité : Sur des graphes de taille Web avec plus de 30 milliards d'arêtes, toutes les bases de référence CPU (y compris la réduction symbolique multicœur) soit ont expiré, soit ont manqué de mémoire. La méthode proposée a réussi à calculer les colorations stables pour ces instances, atteignant des réductions de ~39–46 % dans le nombre de blocs.

Signification et Revendications
L'article affirme que ce travail est le premier à réussir le calcul de colorations stables 1-WL sur des graphes de taille Web avec plus de 30 milliards d'arêtes, une tâche auparavant irréalisable en raison des limitations de mémoire et de calcul. La signification réside dans :

• Permettre l'Analyse de l'Expressivité : Permettre l'analyse empirique de l'expressivité des GNN sur des jeux de données réels et massifs plutôt que de se limiter à de petits ou moyens graphes.
• Adéquation Architecturale : Démontrer que le problème 1-WL, traditionnellement perçu comme séquentiel et global, peut être efficacement parallélisé et distribué en utilisant des formulations algébriques linéaires et des stratégies de batching adaptées aux architectures GPU.
• Robustesse : Fournir une méthode qui est à la fois numériquement stable (via l'arithmétique entière) et scalable en mémoire (via le batching), surmontant les modes de défaillance spécifiques des approches antérieures.

Les auteurs concluent que leur méthode permet l'analyse de l'expressivité 1-WL à des échelles auparavant inaccessibles, tout en notant que la correspondance entre coloration stable et partitions équitables offre des voies futures potentielles pour appliquer ces techniques aux tâches de Machine Learning sur graphes en aval et à la réduction de modèles pour les systèmes dynamiques non linéaires.