Generalizing GNNs with Tokenized Mixture of Experts

Sintesi Tecnica: Generalizzazione di GNN tramite Mixture of Experts Tokenizzati

1. Formulazione del Problema: Il "Triangolo Impossibile" della Distribuzione Congelata

Il documento affronta un limite critico nella distribuzione delle Graph Neural Networks (GNN). In ambienti di produzione, le GNN vengono tipicamente distribuite come "istantanee congelate" (frozen snapshots) senza aggiornamenti dei parametri tra un rilascio e l'altro. Durante questa finestra di congelamento, un singolo modello fisso deve soddisfare simultaneamente tre obiettivi contrastanti, definiti il "triangolo impossibile":

1. Fit (Adattamento): Elevate prestazioni su dati puliti, in-distribuzione.
2. Generalizzazione: Robustezza sotto shift di distribuzione (es. cambiamenti nel grado dei nodi o nell'omofilia).
3. Stabilità: Resistenza alle perturbazioni degli input (es. masking delle feature o cancellazione di archi).

Gli autori sostengono che le GNN standard falliscano nel bilanciare questi obiettivi perché impiegano una singola regola di inferenza fissa (inferenza statica). Teoricamente, imporre la stabilità (insensibilità alle perturbazioni) richiede spesso la soppressione dei segnali ad alta frequenza, che sono frequentemente rilevanti per il fit e per la generalizzazione a ambienti diversificati. Al contrario, ottimizzare per fit e generalizzazione può rendere il modello sensibile al rumore.

Sebbene la Computazione Condizionale all'Istanza (ICC) — che instrada diversi input attraverso meccanismi differenti — offra una potenziale soluzione espandendo la famiglia di percorsi computazionali disponibili, i metodi ICC esistenti introducono nuove fragilità: le decisioni di routing possono derivare sotto shift di distribuzione, e le fluttuazioni del routing possono amplificare la sensibilità degli strati successivi alle perturbazioni.

2. Metodologia: STEM-GNN

Per risolvere la tensione tra copertura (generalizzazione) e stabilità, gli autori propongono STEM-GNN (Stable TokEnized Mixture-of-Experts GNN). Si tratta di un framework di pre-training seguito da fine-tuning progettato per operazionalizzare una ICC robusta attraverso tre componenti strettamente accoppiati:

A. MoE Encoder per l'Espansione della Copertura
Invece di una singola regola di message-passing, STEM-GNN utilizza un encoder Mixture-of-Experts (MoE).

• Meccanismo: Ogni nodo è instradato attraverso una combinazione dipendente dall'input di operatori esperti condivisi tramite un router leggero.
• Obiettivo: Questo permette a un singolo set di parametri congelati di realizzare una famiglia diversificata di meccanismi efficaci, migliorando la copertura attraverso condizioni di test eterogenee (es. grafi con gradi o livelli di omofilia variabili).

B. VQ Tokenization per la Stabilizzazione della Rappresentazione
Per evitare che l'instabilità delle decisioni di routing e la deriva della rappresentazione si propagino alla testa di predizione, gli autori introducono un'interfaccia di tokenizzazione vettoriale (VQ).

• Meccanismo: Gli output continui dell'encoder vengono discretizzati in un codebook fisso tramite l'appaiamento ai vicini più prossimi prima di entrare nella testa di predizione.
• Obiettivo: Piccole perturbazioni nell'output dell'encoder (causate da rumore di input o fluttuazioni di routing) che non superano i confini di quantizzazione risultano in una variazione nulla dell'input della testa di predizione. Ciò crea un'interfaccia stabile che assorbe la minore deriva della rappresentazione.

C. Lipschitz-Regularized Prediction Head per il Controllo della Sensibilità
Per limitare l'impatto delle variazioni rimanenti (es. quando avvengono switch di token), la testa di predizione è regolarizzata.

• Meccanismo: Una penalità sulla norma di Frobenius viene applicata ai pesi della testa di predizione lineare, limitando la sua costante di Lipschitz.
• Obiettivo: Questo limita l'amplificazione di qualsiasi variazione residua dell'input (inclusi gli switch discreti dei token) nell'output finale, garantendo che il modello non reagisca eccessivamente ai cambiamenti a monte.

3. Contributi Teorici

Il documento fornisce un'analisi teorica formale che confronta l'Inferenza Statica con la Computazione Condizionale all'Istanza (ICC):

• Tensione dell'Inferenza Statica: Gli autori dimostrano che, sotto un vincolo di stabilità di Lipschitz, l'inferenza statica affronta una "tensione del testimone" (witness tension) strutturale. Per mantenere la stabilità, il modello deve limitare la sua dipendenza dai componenti ad alta frequenza, il che induce un limite inferiore non nullo (un "pavimento") sul rischio nel peggior scenario ambientale.
• Decomposizione del Rischio ICC: Per l'ICC, gli autori decompongono i rischi in due leve distinte:
  1. Copertura vs Selezione: Il rischio OOD è separato tra quanto bene i meccanismi disponibili coprono condizioni diverse e quanto accuratamente il router seleziona tra di essi.
  2. Sensibilità vs Deriva: Il rischio di stabilità è separato tra la sensibilità di un percorso fisso e l'amplificazione causata dalla deriva del routing.
     Questa analisi giustifica il design di STEM-GNN: MoE affronta la copertura, VQ affronta la deriva/stabilità e la regolarizzazione di Lipschitz affronta l'amplificazione della sensibilità.

4. Risultati Sperimentali

Gli autori valutano STEM-GNN su otto benchmark che spaziano dalla classificazione di nodi, alla predizione di link e alla classificazione di grafi (inclusi Cora, PubMed, Wiki-CS, WN18RR, FB15K-237, HIV e PCBA).

• Prestazioni Clean: STEM-GNN raggiunge la migliore accuratezza media (80,79%) su tutti i benchmark, superando solidi baseline tra cui GFT, GIANT e GraphMAE. Dimostra una ridotta sensibilità ai seed casuali rispetto ad altri metodi.
• Robustezza alle Perturbazioni: Sotto perturbazioni al tempo di inferenza (feature masking ed edge deletion), STEM-GGN degrada più gradualmente rispetto ai baseline, mantenendo un'accuratezza superiore all'aumentare della severità del rumore.
• Generalizzazione OOD: Sotto shift di grado e shift di omofilia, STEM-GNN raggiunge le migliori prestazioni medie nei bucket OOD, particolarmente in scenari impegnativi di alto grado o alta omofilia.
• Bilanciamento Tri-Obiettivo: In una valutazione unificata di Fit, OOD-worst e Perturb-mean, STEM-GNN ottiene il miglior punteggio medio su tutti i benchmark di nodi. Fondamentalmente, migliora tutti e tre gli assi simultaneamente rispetto ai baseline, invece di scambiarne uno per l'altro.
• Studi di Ablazione: Rimuovere uno qualsiasi dei tre componenti (MoE, VQ o regolarizzazione di Lipschitz) porta a un calo della media tri-obiettivo, confermando che i componenti sono complementari piuttosto che ridondanti.

5. Significato e Rivendicazioni

Il documento sostiene che STEM-GNN riesce a inquadrare con successo la generalizzazione robusta dei grafi sotto distribuzione congelata come una tensione tri-obiettivo governata dal meccanismo di computazione al tempo di test, non solo dall'obiettivo di addestramento.

Le rivendicazioni chiave includono:

• Framework Unificato: È il primo metodo ad affrontare congiuntamente l'espansione della copertura dei meccanismi e il controllo esplicito della stabilità al tempo di distribuzione (tramite tokenizzazione e vincoli di Lipschitz) nell'apprendimento sui grafi.
• Insight Teorico: Fornisce una prova formale che l'inferenza statica affronta intrinsecamente un compromesso tra stabilità e generalizzazione, mentre l'ICC offre nuovi leve di design per mitigare questa tensione.
• Efficacia Pratica: Il framework mantiene forti prestazioni "clean" pur migliorando significativamente il bilanciamento tra robustezza agli shift di distribuzione e alle perturbazioni, rendendolo adatto a scenari reali di distribuzione congelata dove i modelli non possono essere aggiornati al tempo di test.

Gli autori osservano che l'approccio si basa su un vincolo di stabilità di Lipschitz e su una famiglia finita di meccanismi instradati, riconoscendo che il rilassamento di questi assunti è una direzione per il lavoro futuro. Il codice e i dati sono resi disponibili per facilitare la riproducibilità.