Des étoiles aux sous-graphes : améliorer tout GNN grâce à une prise de conscience de la structure locale

Les Réseaux de Neurones à Passage de Messages (MPNN) constituent un type courant de Réseau de Neurones sur Graphes (GNN), dans lesquels la représentation de chaque nœud est calculée de manière récursive en agrégant les représentations (messages) de ses voisins immédiats, selon un schéma semblant une étoile. Les MPNN sont particulièrement attrayants en raison de leur efficacité et de leur capacité à échelle, mais leur expressivité est bornée supérieurement par le test d’isomorphisme de Weisfeiler-Lehman d’ordre 1 (1-WL). À cette fin, les travaux antérieurs ont proposé des modèles hautement expressifs au prix d’une perte de scalabilité et parfois d’une performance de généralisation réduite. Notre travail se situe entre ces deux extrêmes : nous introduisons un cadre général permettant d’améliorer significativement l’expressivité de tout MPNN, avec un surcoût limité en scalabilité et une performance pratique considérablement accrue. Cela est réalisé en étendant l’agrégation locale des MPNN, initialement restreinte aux motifs d’étoile, à des motifs de sous-graphes généraux (par exemple, les k-égonets) : dans notre cadre, la représentation de chaque nœud est calculée comme l’encodage d’un sous-graphe induit environnant, plutôt que comme l’encodage uniquement des voisins immédiats (c’est-à-dire une étoile). Nous choisissons l’encodeur de sous-graphe comme étant un GNN (principalement des MPNN, afin de préserver la scalabilité), afin de concevoir un cadre général pouvant servir de wrapper pour améliorer tout GNN. Nous appelons notre méthode GNN-AK (GNN en tant que noyau), car ce cadre ressemble à un réseau de neurones convolutif en remplaçant le noyau par des GNN. Théoriquement, nous démontrons que notre cadre est strictement plus puissant que le 1-WL et le 2-WL, et n’est pas moins puissant que le 3-WL. Nous avons également conçu des stratégies d’échantillonnage de sous-graphes qui réduisent considérablement la charge mémoire et améliorent la vitesse, tout en maintenant une performance élevée. Notre méthode établit de nouveaux états de l’art, avec des écarts significatifs, sur plusieurs tâches classiques d’apprentissage automatique sur graphes ; plus précisément, un MAE de 0,08 sur ZINC, 74,79 % et 86,887 % de précision sur CIFAR10 et PATTERN respectivement.