Implicit Regularization of Mini-Batch Training in Graph Neural Networks

Technische Samenvatting: Impliciete Regularisatie van Mini-Batch Training in Graph Neural Networks

Probleemstelling
Het trainen van Graph Neural Networks (GNN's) op grote schaal grafen vereist doorgaans mini-batch training om geheugen- en rekenkundige beperkingen te managen. In tegenstelling tot training op onafhankelijke en identiek verdeelde (i.i.d.) data, is mini-batching in GNN's fundamenteel anders: het samplen van een subgraaf verandert de grafentopologie, introduceert randeffecten en wijzigt de interacties die beschikbaar zijn voor het model. Bijgevolg is de keuze van de sampler niet louter een implementatiedetail, maar bepaalt deze de stochastische gradiënten en het effectieve leervermogen. Vorig werk heeft zich gericht op het ontwikkelen van "structuurbewuste" samplers (bijv. Neighborhood Sampling, ClusterGCN, GraphSAINT) die lokale connectiviteit behouden om de variatie van embeddings te verminderen. Het theoretische inzicht in welk doelwit er daadwerkelijk wordt geoptimaliseerd bij training op deze gesampelde subgrafen blijft echter onvolledig. Specifiek is het onduidelijk hoe de samplingstrategie de impliciete regularisatie van het optimalisatieproces beïnvloedt.

Methodologie en Theoretisch Kader
De auteurs onderzoeken Random Node Sampling (RNS), een eenvoudig schema waarbij knopen uniform willekeurig worden gesampled en training doorgaat op de geïnduceerde subgraaf. Hoewel RNS lokale structuur verwaarloost en "onvolledige" omgevingen creëert in vergelijking met de volledige graaf, demonstreert het artikel de verrassende effectiviteit ervan.

Om dit te verklaren, passen de auteurs backward error analysis toe op graf mini-batch Stochastic Gradient Descent (SGD). Deze techniek interpreteert de discrete optimalisatietraject niet als een benadering van een continue gradiëntstroom, maar als de exacte stroom van een gewijzigd systeem.

• Gewijzigd Doelwit: De analyse bewijst dat training op gesampelde subgrafen impliciet een gewijzigde verliesfunctie minimaliseert:
  $$\tilde{L}_{\text{Sampl}}(w) = \bar{L}_{\text{Sampl}}(w) + \frac{\epsilon}{4} \|\nabla \bar{L}_{\text{Sampl}}(w)\|^2 + R(w)$$
  waarbij $\bar{L}_{\text{Sampl}}$ het epoch-gemiddelde gesamplede verlies is, en $R(w)$ een regularisator is die evenredig is met de variatie van mini-batch gradiënten.
• Sampler-Afhankelijkheid: De term $R(w)$ wordt direct gevormd door de sampler. De auteurs tonen aan dat RNS, hoewel het grafstructuur negeert, mini-batches produceert waarvan het verwachte verlies dichter bij het verlies van de volledige graaf ligt (lage bias) en waarvan de gradiënten per batch aanzienlijk lagere variatie vertonen in vergelijking met structuurbewuste alternatieven.
• Theoretisch Voordeel: In tegenstelling tot op structuur gebaseerde samplers (bijv. ClusterGCN, GraphSAINT), die bias introduceren door niet-uniforme knoopselectie of vervorming van de graad, behoudt RNS de power-law graadverdeling van de oorspronkelijke graaf in elke mini-batch. Deze structurele behoud correleert met een regularisatieterm voor gradiëntvariatie die orde van grootte kleiner is dan die van andere methoden, wat leidt tot een meer stabiel impliciet doelwit.

Belangrijkste Bijdragen

1. Theoretische Analyse: Het artikel biedt een backward error analyse voor GNN mini-batch training, waarbij wordt bewezen dat de sampler het doelwit impliciet regulariseert via een straal voor gradiëntvariatie. Het stelt vast dat RNS een goed gecontroleerde bias induceert ten opzichte van het verlies van de volledige graaf en een minimale variatie-regularisatieterm.
2. Empirische Benchmarking: De auteurs benchmarken vijf samplingstrategieën (RNS, Neighborhood Sampling, ClusterGCN, GraphSAINT, LADIES) over drie grote datasets (OGBN-ARXIV, OGBN-PRODUCTS, POKEC) en vier GNN-architecturen (GCN, GraphSAGE, GAT, SGFormer).
3. Praktische Impact: RNS wordt gepresenteerd als een "drop-in" vervanging voor training op de volledige graaf, waarbij slechts één geheel getal hyperparameter nodig is ($m$, het aantal batches).

Resultaten

• Prestaties: RNS komt overeen met of overtreft training op de volledige graaf op 8 van de 10 benchmarks (over variërende grafmaten, dichtheden en homofilie-niveaus). Op OGBN-PRODUCTS bereikt RNS 83,11% nauwkeurigheid tegenover 80,42% voor training op de volledige graaf. Op OGBN-ARXIV en POKEC overtreft RNS ook de baselines voor training op de volledige graaf.
• Efficiëntie: RNS bereikt 2× tot 12× snelheidswinst in wandkloktijd in vergelijking met training op de volledige graaf en verlaagt het piekgebruik van GPU-geheugen met maximaal 3×. Dit wordt toegeschreven aan het verwerken van minder randen per epoch (ongeveer $|E|/m$) en het uitvoeren van $m$ optimalisatorstappen per epoch.
• Generalisatie: De methode generaliseert over verschillende message-passing-architecturen (GCN, GraphSAGE) en diverse datasets van de GraphLand-benchmark. Het toont robuustheid in zowel schaarste (Random Low) als overvloedige (Random High) toezichtregimes.
• Ablatie: Experimenten bevestigen dat de prestatiewinst afhankelijk is van zowel de samplingbias als de gekoppelde mini-batch updates (de $R(w)$-term). Het verwijderen van het koppelende effect (gradiënten accumuleren voordat er wordt bijgewerkt) vermindert het prestatievoordeel, wat de theoretische afleiding valideert.

Betekenis en Claims
Het artikel claimt dat RNS moet worden beschouwd als een sterke, theoretisch onderbouwde standaard voor schaalbare transductieve GNN-training. De betekenis ligt in het herformuleren van de keuze van de grafsampler als een vorm van impliciete regularisatie. De auteurs betogen dat de eenvoud van RNS geen beperking is, maar een eigenschap: door complexe structuurbehoudende heuristieken te vermijden, minimaliseert het per ongeluk de variatie van de impliciete regularisator, wat leidt tot een meer stabiel en effectief optimalisatielandschap.

De auteurs blijven bescheiden over bepaalde aspecten:

• Adam Optimizer: De backward error analyse is afgeleid voor vanilla SGD. Hoewel empirische resultaten tonen dat RNS goed presteert met Adam, blijft de formele uitbreiding van de theorie van impliciete regularisatie naar adaptieve optimalizers een open vraag vanwege de nonlineariteit die wordt geïntroduceerd door normalisatie per coördinaat.
• Gevoeligheid voor Hyperparameters: Het aantal batches per epoch ($m$) wordt geïdentificeerd als een kritieke hyperparameter. Hoewel RNS eenvoudiger is om te tunen dan methoden met meerdere samplingparameters (fanout, wandlengte), moet $m$ nog steeds zorgvuldig worden geselecteerd, aangezien de nauwkeurigheid gevoelig kan zijn voor de waarde ervan, met name met betrekking tot de subgraafdiameter.
• Architectuurlimieten: De methode blijkt effectief te zijn voor standaard message-passing GNN's. Echter, voor modellen met globale attentiemechanismen (zoals SGFormer) kan RNS onderpresteren als de globale interacties worden afgekapt door het random node sampling, wat suggereert dat de methode niet universeel architectuur-agnostisch is zonder verder onderzoek.

Concluderend stelt het artikel vast dat de "eenvoudigste" samplingstrategie, vaak van de hand gewezen omdat deze lokale structuur verwaarloost, eigenlijk een superieur impliciet doelwit biedt voor GNN-training door gradiëntvariatie te minimaliseren, en een praktische en efficiënt alternatief biedt voor complexe structuurbewuste samplingmethoden.