Implicit Regularization of Mini-Batch Training in Graph Neural Networks

Résumé technique : Régularisation implicite de l'entraînement par mini-lots dans les réseaux de neurones à graphes

Énoncé du problème
L'entraînement de réseaux de neurones à graphes (GNN) sur des graphes à grande échelle nécessite généralement un entraînement par mini-lots pour gérer les contraintes de mémoire et de calcul. Contrairement à l'entraînement sur des données indépendantes et identiquement distribuées (i.i.d.), le mini-lotage dans les GNN est fondamentalement distinct : l'échantillonnage d'un sous-graphe modifie la topologie du graphe, introduit des effets de bord et change les interactions disponibles pour le modèle. Par conséquent, le choix de l'échantillonneur n'est pas un simple détail d'implémentation, mais définit les gradients stochastiques et le problème d'apprentissage effectif. Les travaux antérieurs se sont concentrés sur le développement d'échantillonneurs « conscients de la structure » (par exemple, échantillonnage de voisinage, ClusterGCN, GraphSAINT) qui préservent la connectivité locale pour réduire la variance des plongements. Cependant, la compréhension théorique de l'objectif réellement optimisé lors de l'entraînement sur ces sous-graphes échantillonnés reste incomplète. Plus précisément, il est unclear comment la stratégie d'échantillonnage influence la régularisation implicite du processus d'optimisation.

Méthodologie et cadre théorique
Les auteurs étudient l'échantillonnage aléatoire de nœuds (RNS), un schéma simple où les nœuds sont échantillonnés uniformément au hasard, et l'entraînement se poursuit sur le sous-graphe induit. Bien que le RNS rejette la structure locale et crée des voisinages « incomplets » par rapport au graphe complet, l'article démontre son efficacité surprenante.

Pour l'expliquer, les auteurs appliquent l'analyse d'erreur rétroactive à la descente de gradient stochastique (SGD) par mini-lots sur graphes. Cette technique interprète la trajectoire d'optimisation discrète non pas comme une approximation d'un flot de gradient continu, mais comme le flot exact d'un système modifié.

• Objectif modifié : L'analyse prouve que l'entraînement sur des sous-graphes échantillonnés minimise implicitement une fonction de perte modifiée :
  $$\tilde{L}_{\text{Sampl}}(w) = \bar{L}_{\text{Sampl}}(w) + \frac{\epsilon}{4} \|\nabla \bar{L}_{\text{Sampl}}(w)\|^2 + R(w)$$
  où $\bar{L}_{\text{Sampl}}$ est la perte échantillonnée moyennée sur l'époque, et $R(w)$ est un régularisateur proportionnel à la variance des gradients de mini-lots.
• Dépendance à l'échantillonneur : Le terme $R(w)$ est directement façonné par l'échantillonneur. Les auteurs montrent que, bien que le RNS ignore la structure du graphe, il produit des mini-lots dont la perte attendue est plus proche de la perte du graphe complet (faible biais) et dont les gradients par lot présentent une variance significativement plus faible par rapport aux alternatives conscientes de la structure.
• Avantage théorique : Contrairement aux échantillonneurs basés sur la structure (par exemple, ClusterGCN, GraphSAINT), qui introduisent un biais par une sélection de nœuds non uniforme ou une distorsion des degrés, le RNS maintient la distribution des degrés en loi de puissance du graphe original dans chaque mini-lot. Cette préservation structurelle est corrélée à un terme de régularisation de variance de gradient qui est d'ordres de grandeur plus petit que celui d'autres méthodes, conduisant à un objectif implicite plus stable.

Contributions clés

1. Analyse théorique : L'article fournit une analyse d'erreur rétroactive pour l'entraînement par mini-lots des GNN, prouvant que l'échantillonneur régularise implicitement l'objectif via une pénalité de variance de gradient. Il établit que le RNS induit un biais bien contrôlé par rapport à la perte du graphe complet et un terme de régularisation de variance minimal.
2. Étalonnage empirique : Les auteurs évaluent cinq stratégies d'échantillonnage (RNS, échantillonnage de voisinage, ClusterGCN, GraphSAINT, LADIES) sur trois jeux de données à grande échelle (OGBN-ARXIV, OGBN-PRODUCTS, POKEC) et quatre architectures de GNN (GCN, GraphSAGE, GAT, SGFormer).
3. Impact pratique : Le RNS est présenté comme un remplacement « plug-and-play » de l'entraînement sur graphe complet, ne nécessitant qu'un seul hyperparamètre entier ($m$, le nombre de lots).

Résultats

• Performance : Le RNS égale ou dépasse l'entraînement sur graphe complet sur 8 benchmarks sur 10 (à travers différentes tailles de graphes, densités et niveaux d'homophilie). Sur OGBN-PRODUCTS, le RNS atteint une précision de 83,11 % contre 80,42 % pour l'entraînement sur graphe complet. Sur OGBN-ARXIV et POKEC, le RNS dépasse également les bases de référence sur graphe complet.
• Efficacité : Le RNS réalise des accélérations de temps réel de 2 à 12 fois par rapport à l'entraînement sur graphe complet et réduit l'utilisation maximale de la mémoire GPU jusqu'à 3 fois. Cela est attribué au traitement de moins d'arêtes par époque (environ $|E|/m$) et à l'exécution de $m$ étapes d'optimiseur par époque.
• Généralisation : La méthode généralise à différentes architectures de passage de messages (GCN, GraphSAGE) et à des jeux de données divers du benchmark GraphLand. Elle montre une robustesse à la fois dans des régimes de supervision rares (Random Low) et abondants (Random High).
• Ablation : Les expériences confirment que le gain de performance repose à la fois sur le biais d'échantillonnage et les mises à jour en chaîne des mini-lots (le terme $R(w)$). Supprimer l'effet de chaîne (accumuler les gradients avant la mise à jour) diminue l'avantage de performance, validant la dérivation théorique.

Signification et affirmations
L'article affirme que le RNS devrait être considéré comme une valeur par défaut forte et théoriquement fondée pour l'entraînement transductif évolutif des GNN. La signification réside dans le recadrage du choix de l'échantillonneur de graphe comme une forme de régularisation implicite. Les auteurs soutiennent que la simplicité du RNS n'est pas une limitation mais un atout : en évitant des heuristiques complexes de préservation de la structure, il minimise involontairement la variance du régularisateur implicite, conduisant à un paysage d'optimisation plus stable et efficace.

Les auteurs restent modestes concernant certains aspects :

• Optimiseur Adam : L'analyse d'erreur rétroactive est dérivée pour le SGD classique. Bien que les résultats empiriques montrent que le RNS fonctionne bien avec Adam, l'extension formelle de la théorie de la régularisation implicite aux optimiseurs adaptatifs reste une question ouverte en raison de la non-linéarité introduite par la normalisation coordonnée par coordonnée.
• Sensibilité aux hyperparamètres : Le nombre de lots par époque ($m$) est identifié comme un hyperparamètre critique. Bien que le RNS soit plus simple à régler que les méthodes avec plusieurs paramètres d'échantillonnage (fanout, longueur de marche), $m$ doit toujours être sélectionné avec soin, car la précision peut être sensible à sa valeur, en particulier concernant le diamètre du sous-graphe.
• Limites architecturales : La méthode s'avère efficace pour les GNN standards de passage de messages. Cependant, pour les modèles avec des mécanismes d'attention globale (comme SGFormer), le RNS peut sous-performer si les interactions globales sont tronquées par l'échantillonnage aléatoire de nœuds, suggérant que la méthode n'est pas universellement agnostique à l'architecture sans investigations supplémentaires.

En conclusion, l'article établit que la stratégie d'échantillonnage la plus « simple », souvent rejetée pour avoir écarté la structure locale, fournit en réalité un objectif implicite supérieur pour l'entraînement des GNN en minimisant la variance des gradients, offrant une alternative pratique et efficace aux méthodes complexes d'échantillonnage conscient de la structure.