Invariant-Based Weight Sharing for Message Passing

Sintesi Tecnica: Condivisione dei Pesi Basata su Invarianti per il Passaggio di Messaggi

Enunciato del Problema
Le reti neurali a passaggio di messaggi (MPNN) costituiscono un framework dominante per l'apprendimento di rappresentazioni di dati strutturati a grafo. Tuttavia, le MPNN standard e i transformer per grafi si basano tipicamente su una condivisione dei pesi fondata esclusivamente sulle caratteristiche dei nodi o degli archi, oppure apprendono coefficienti di attenzione specifici per l'istanza. Questo approccio limita la capacità del modello di catturare pattern strutturali intrinseci, poiché i pesi non sono esplicitamente vincolati alla topologia del grafo. Inoltre, le MPNN standard sono spesso limitate all'aggregazione del vicinato locale, rendendole suscettibili all'over-smoothing e all'over-squashing in architetture profonde, mentre i transformer per grafi spesso mancano di un riutilizzo sistematico dei parametri tra sottografi strutturalmente equivalenti.

Metodologia: ShareGNN
Gli autori introducono gli ShareGNN, una nuova architettura fondata sulla condivisione dei pesi basata su invarianti. Il principio fondamentale consiste nell'indicizzare i parametri apprendibili direttamente tramite invarianti del grafo—funzioni preservate sotto permutazioni dei nodi—piuttosto che tramite dimensioni fisse delle caratteristiche o posizioni specifiche per l'istanza.

1. Condivisione dei Pesi Basata su Invarianti:

   • Codificatore (Passaggio di Messaggi): Invece di matrici di pesi fisse, il modello costruisce dinamicamente i pesi per il passaggio di messaggi per ogni coppia di nodi $(v, w)$ basandosi su una terna invariante: $(l(v), l(w), d(v, w))$, dove $l$ è una funzione di etichettatura (ad esempio, numero atomico, etichetta di Weisfeiler-Leman, conteggio di pattern) e $d$ è la distanza del percorso più breve. Tutte le coppie di nodi che condividono la stessa terna invariante condividono un singolo peso scalare apprendibile $\omega_\tau$. Ciò consente il flusso di informazioni tra nodi arbitrari in un singolo strato, non solo tra vicini immediati.
   • Decodificatore (Pooling): Analogamente, i termini di bias e i pesi di pooling per l'aggregazione a livello di grafo sono indicizzati dalle etichette invarianti dei singoli nodi.
   • Costruzione Dinamica: Per qualsiasi grafo in ingresso, le matrici di pesi pertinenti sono assemblate al volo da un pool di parametri condiviso. Ciò garantisce che il modello sia equivariante rispetto alle permutazioni nel codificatore e invariante rispetto alle permutazioni nel decodificatore per costruzione.

2. Funzioni di Etichettatura:
   L'espressività del modello è controllata dalla scelta delle funzioni di etichettatura $l$. Il lavoro esplora:

   • Etichettatura Molecolare: Numeri atomici.
   • Etichettatura Weisfeiler-Leman (WL): Raffinamento iterativo delle etichette dei nodi basato sui multinsiemi dei vicini.
   • Etichettatura dei Pattern: Conteggi di specifici pattern di sottografi (ad esempio, triangoli, cicli) che coinvolgono un nodo.
     Questi possono essere combinati per catturare informazioni strutturali a più livelli.

3. Architettura:
   Gli ShareGNN utilizzano una semplice struttura codificatore-decodificatore. Il codificatore propaga le informazioni attraverso distanze arbitrarie in un singolo strato, e il decodificatore aggrega gli embedding dei nodi in una rappresentazione del grafo utilizzando un pooling indicizzato da invarianti. L'architettura supporta estensioni multi-testa in cui diverse teste utilizzano diversi insiemi di invarianti.

Principali Contributi

1. Principio di Condivisione dei Pesi Basato su Invarianti: Un nuovo meccanismo per la condivisione dei pesi che incorpora esplicitamente la struttura del grafo, consentendo un riutilizzo sistematico dei parametri tra coppie di nodi strutturalmente equivalenti all'interno e tra grafi diversi.
2. Architettura ShareGNN: Un'istanza pratica di questo principio che combina il passaggio di messaggi con una connettività simile a quella dei transformer (interazioni tutti-con-tutti), ma con parametri vincolati a invarianti strutturali piuttosto che a punteggi di attenzione appresi.
3. Validazione Teorica ed Empirica:
   • Espressività: Gli autori dimostrano che gli ShareGNN sono almeno tanto espressivi quanto la funzione di etichettatura scelta. Se due grafi sono distinguibili dalla funzione di etichettatura (ad esempio, etichette WL), lo ShareGNN può distinguerli.
   • Controllo: Il modello offre un controllo esplicito sulla complessità e sull'espressività attraverso la selezione degli invarianti.
   • Prestazioni: Gli esperimenti dimostrano che gli ShareGNN superano le MPNN standard (ad esempio, GIN, GAT) e sono competitivi con i transformer per grafi su vari benchmark, inclusa la classificazione molecolare, l'analisi di reti sociali e il conteggio di sottstrutture.

Risultati

• Classificazione di Grafi: Su dataset diversificati (NCI1, NCI109, MUTAG, IMDB-B/M), gli ShareGNN raggiungono una precisione all'avanguardia o altamente competitiva. In particolare, raggiungono una precisione quasi perfetta (100%) su benchmark sintetici (RingTransfer, Snowflakes) progettati per mettere alla prova le MPNN standard, dimostrando la loro capacità di catturare dipendenze a lungo raggio e di distinguere grafi indistinguibili dal test 1-WL.
• Conteggio di Sottstrutture: Sul Benchmark per il Conteggio di Sottstrutture, gli ShareGNN raggiungono risultati all'avanguardia in impostazioni a compito singolo e rimangono competitivi in impostazioni a più compiti, validando il fatto che invarianti espressivi portino a modelli espressivi anche con architetture superficiali (singolo strato di codificatore).
• Scalabilità: Il modello scala su dataset di grandi dimensioni come ZINC (12k e 250k grafi) mantenendo un'architettura superficiale, evitando i problemi di over-smoothing comuni nelle MPNN profonde.
• Efficienza: Studi di ablazione mostrano che aumentare la massima distanza di passaggio dei messaggi $D$ ha un impatto più significativo sulle prestazioni rispetto all'aggiunta di ulteriori strati di codificatore. Inoltre, la potatura delle terne invarianti rare riduce significativamente il numero di parametri senza sacrificare l'accuratezza.

Significato e Affermazioni
Il lavoro afferma che la condivisione dei pesi basata su invarianti fornisce un'alternativa fondata alla condivisione dei pesi basata su caratteristiche e all'attenzione specifica per l'istanza. Vincolando i pesi a invarianti strutturali, gli ShareGNN offrono:

• Migliore Interpretabilità: I pesi appresi possono essere interpretati direttamente in termini di pattern strutturali (ad esempio, distanze specifiche tra specifici tipi di atomi).
• Trasferibilità: I parametri sono naturalmente trasferibili tra grafi con pattern strutturali simili, suggerendo potenzialità per modelli fondazionali di grafi.
• Controllo dell'Espressività: Il modello consente ai ricercatori di controllare esplicitamente il potere discriminativo della rete selezionando invarianti appropriati, colmando il divario tra espressività teorica (ad esempio, test k-WL) e implementazione pratica.

Gli autori riconoscono le limitazioni, inclusi i costi di pre-elaborazione per il calcolo degli invarianti e delle distanze pairwise, e la complessità di memoria quadratica nel caso peggiore, sebbene notino che ciò è paragonabile ai transformer per grafi e può essere mitigato mediante potatura. Osservano inoltre che l'approccio attuale richiede invarianti definiti dall'utente, lasciando spazio a futuri lavori sulla selezione automatizzata degli invarianti.