Technische Samenvatting: Early-Exit Graph Neural Networks voor Linkvoorspelling

Probleemstelling

Graph Neural Networks (GNN's) zijn de standaard voor het leren op graafgestructureerde data, maar hun toepassing in grootschalige domeinen wordt vaak beperkt door inferentielatentie. Twee fundamentele problemen hinderen de efficiëntie en prestaties:

1. Over-smoothing: Naarmate de diepte van een GNN toeneemt, neigen node-embeddings naar het convergeren naar onderscheidingsloze vectoren, wat de effectieve diepte van de meeste architecturen beperkt tot ongeveer 20 lagen.
2. Niet-adaptiviteit: Standaard GNN's voeren een vast aantal message-passing stappen uit voor elke node, ongeacht de lokale moeilijkheidsgraad. In grafen met heterogene substructuren (bijv. dichte versus ijle regio's), resulteert dit in verspilde berekeningen voor nodes waar voorspellingen met minder lagen gemaakt hadden kunnen worden.

Hoewel early-exit strategieën bestaan voor node- en graafclassificatie, vertrouwen deze doorgaans op hulp-budgetverliezen (auxiliary budget losses) om het stoppen af te dwingen. Bovendien heeft geen enkel eerder werk succesvol auxiliary-loss-vrije early exiting toegepast op linkvoorspelling, een taak die gecompliceerd wordt door het feit dat voorspellingen afhangen van paren node-embeddings en variërende subgraafstructuren.

Methodologie

De auteurs stellen twee early-exit strategieën voor linkvoorspelling voor die opereren zonder hulp-verliezen, waardoor het exit-beleid impliciet kan voortkomen uit de taak-objective (linkvoorspellingsverlies) alleen. De aanpak is geïnspireerd door het EEGNN-framework [1], waarbij gebruik wordt gemaakt van stabiele Neural ODE backbones (specifiek SAS-GNN) en standaard GCN's.

Kernmechanismen

1. Gumbel-Softmax Exit Mechanisme:
   Bij elke laag $l$ en voor elke node $u$ voorspelt een lichtgewicht confidence head een waarschijnlijkheidsverdeling over acties $\{CONTINUE, EXIT\}$. Een temperatuurparameter $\nu$ regelt de stochasticiteit, en de Gumbel-Softmax truc wordt gebruikt om differentiabele sampling mogelijk te maken. Om stabiliteit te garanderen, wordt de temperatuur van onderaf afgeklempt (bijv. $\nu \ge 0.1$) om gradiëntdivergentie te voorkomen.

2. Node-gebaseerde Early Exiting:
   Elke node beslist onafhankelijk of hij bij elke laag stopt. Wanneer een node stopt, wordt de embedding bevroren ($h^{l+1}_u = h^l_u$), maar de node blijft deelnemen aan message passing voor zijn buren. De uiteindelijke embedding die wordt gebruikt voor linkvoorspelling is de representatie op de laag waar de node voor het eerst besloot te stoppen.

3. Subgraph-gebaseerde Early Exiting:
   Om inconsistenties aan te pakken waarbij een node bevriest terwijl zijn buren doorgaan met updaten (of vice versa), introduceert deze methode een buurtconsensusprotocol:

   • Fase 1 (Soft Exit): Nodes genereren een soft exit score op basis van hun eigen vertrouwen.
   • Fase 2 (Hard Exit): Een node mag pas "hard exit" (zijn embedding bevriezen) als zijn lokale buurt ook een hoge readiness staat deelt. De exit-waarschijnlijkheid wordt gegate door de gemiddelde soft-exit staat van de buren van de node. Dit dwingt dichte subgrafen om gezamenlijk te stoppen, wat fungeert als een structurele regularisator.

Theoretische Berekening

De auteurs benaderen de rekenkosten met een formule die rekening houdt met actieve nodes en edges. Ze merken op dat hoewel de FLOP-besparingen aanzienlijk zijn, het realiseren van wall-clock snelheidswinsten op GPU's ondersteuning voor sparse-executie of aangepaste kernels vereist, aangezien naïeve dichte implementaties nog steeds tensorrijen alloceren voor bevroren nodes.

Belangrijkste Bijdragen

1. Nieuwe Strategieën: Het voorstel van node-niveau en subgraph-niveau early-exit strategieën die specifiek zijn aangepast aan linkvoorspellings-taken.
2. Auxiliary-Loss-Vrije Training: Het aantonen dat early exiting natuurlijk kan voortkomen uit de taak-loss alleen, zonder de noodzaak voor auxiliary budget losses, volgend naar het EEGNN-framework.
3. Benchmark Adaptatie: Het aanpassen van deze methoden om te opereren op de HeaRT benchmark, een standaard voor het evalueren van GNN-efficiëntie.
4. Verschuiving van de Pareto Front: Empirisch bewijs dat deze methoden de prestatie-rekenkracht Pareto front verschuiven, waarbij ze hogere MRR en Hit@ metrieken bereiken voor kleinere rekenbudgetten vergeleken met fixed-depth baselines.

Resultaten

De methoden werden geëvalueerd op drie HeaRT datasets: Citeseer, PubMed, en ogbl-ddi, met behulp van GCN en SAS-GNN backbones.

• Exit Diepte Distributie:
  • Op Citeseer (een kleine, dichte citatie-netwerk), stoppen nodes bijna onmiddellijk (vaak bij laag 1), wat suggereert dat de taak minimaal message passing vereist.
  • Op PubMed, vertonen diepere netwerken een duidelijke "bump" in de exit-waarschijnlijkheid rond laag 4, wat aangeeft dat het model leert om een specifieke intrinsieke diepte te benutten voordat het stopt.
• Prestaties vs. Rekenkracht:
  • Rekenbesparingen: Beide strategieën verminderen de cumulatieve FLOPs aanzienlijk vergeleken met fixed-depth baselines. Node-gebaseerde exiting is agressiever en bespaart ongeveer 5$\times$ FLOPs op Citeseer bij laag 15 vergeleken met de baseline.
  • Voorspellingskwaliteit:
    • Citeseer: Subgraph exiting evenaart of overtreft de prestaties van een "oracle" (een fixed-depth model met optimaal gekozen diepte) op Hit@10. Beide methoden presteren aanzienlijk beter dan de baseline.
    • PubMed: De resultaten zijn gemengd; node-gebaseerde exiting presteert slechter dan de baseline op MRR en Hit@1, terwijl subgraph exiting competitief blijft.
    • ogbl-ddi: Beide methoden verbeteren Hit@10 ten opzichte van de baseline, met vergelijkbare MRR.
• Stabiliteit: De subgraph-gebaseerde methode bleek robuuster tegen hyperparameter-variaties en trainingsdynamiek dan de node-gebaseerde methode, omdat het fungeert als een structurele regularisator die voorkomt dat het model vervalt in "instant exit" of "never exit" gedrag.

Betekenis en Beperkingen

Het artikel beweert dat de inferentiesnelheid voor GNN's bij linkvoorspellingsproblemen kan worden verbeterd met weinig tot geen verlies in voorspellingskwaliteit, en in sommige gevallen zelfs met winst. De betekenis ligt in het aantonen dat adaptieve rekenkracht haalbaar is voor linkvoorspelling zonder complexe auxiliary losses.

De auteurs behouden echter een bescheiden toon met betrekking tot hun bevindingen:

• Gevoeligheid: Het gedrag van de methode is zeer gevoelig voor hyperparameters en implementatiekeuzes (bijv. learning rate). Kleine veranderingen kunnen leiden tot stochastische uitkomsten waarbij het model ofwel onmiddellijk stopt of nooit stopt.
• Hardwarebeperkingen: De gerapporteerde FLOP-besparingen vertalen zich niet automatisch naar GPU-snelheidswinsten zonder ondersteuning voor sparse-executie.
• Scope: De evaluatie is beperkt tot drie datasets met gereduceerde trainingsbudgetten (10$\times$ minder stappen dan standaardprotocollen) en slechts twee backbone architecturen (GCN en SAS-GNN). Generalisatie naar attention-gebaseerde of heterogene GNN's blijft ongetest.

De auteurs concluderen dat hoewel de node-gebaseerde aanpak theoretische efficiëntie biedt, de subgraph-gebaseerde aanpak een robuustere weg voorwaarts vertegenwoordigt voor real-world toepassingen vanwege de capaciteit om ruimtelijke consensus af te dwingen en training te stabiliseren.