Riepilogo Tecnico: Graph Set Transformer (GST)

Definizione del Problema

I compiti di machine learning che coinvolgono insiemi di grafi — collezioni di grafi non ordinati dove le previsioni dipendono sia dalla struttura locale del grafo che dal contesto globale dell'insieme — pongono una sfida architettonica significativa. Gli approcci esistenti, come DeepSets e Set Transformer, seguono tipicamente una pipeline sequenziale: codificano prima i singoli grafi in embedding vettoriali fissi utilizzando una Rete Neurale su Grafo (GNN) e successivamente applicano una testa di insieme invariante per permutazione per aggregare questi embedding.

Gli autori sostengono che questa stadiazione sequenziale crei un collo di bottiglia. In tali pipeline, l'informazione a livello di insieme viene ricostruita solo nella fase finale a partire da embedding riassuntivi che sono stati calcolati senza alcuna consapevolezza degli altri grafi nell'insieme. Questa separazione impedisce al modello di ragionare efficacemente su come la struttura locale di un grafo specifico interagisca con il contesto globale dell'intero insieme, una necessità per compiti come la previsione della resa di una reazione o l'identificazione dei centri di reazione, dove il comportamento di un atomo dipende da tutte le molecole partecipanti.

Metodologia

Il documento introduce il Graph Set Transformer (GST), un'architettura neurale progettata per intercalare la propagazione delle caratteristiche a livello di nodo locale con la modellazione del contesto tra i grafi ad ogni livello, anziché scalarli sequenzialmente.

Panoramica dell'Architettura

Il nucleo di GST è il blocco graphsetconv, che elabora un insieme di grafi $S = \{G_1, G_2, \dots\}$. Il blocco opera in cinque passaggi distinti per aggiornare le caratteristiche dei nodi $x_i$:

1. Passaggio di Messaggi Locali: Un tipico strato di Rete Convoluzionale su Grafo (GCN) aggrega le caratteristiche dai vicini a 1 salto (1-hop). Gli autori utilizzano GraphNorm invece di LayerNorm per gestire le dimensioni variabili dei grafi all'interno di un insieme, garantendo la stabilità delle magnitudo di attivazione a livello di grafo.
2. Pooling Attentivo: Ogni grafo nell'insieme viene riassunto in un singolo token $d$-dimensionale ($z_g$) tramite una somma pesata da softmax delle sue caratteristiche nodali. Questo crea un collo di bottiglia per il flusso di informazioni tra i grafi.
3. Contestualizzazione a Livello di Insieme: I token di grafo ($z_g$) di tutti i grafi nell'insieme partecipano all'attenzione reciproca utilizzando un blocco di Multi-Head Self-Attention (MHSA). Questo passaggio cattura il contesto dell'insieme (ad esempio, identificando estremi, riassunti di distribuzione o classifiche) e aggiorna i token di grafo ($\tilde{z}_g$).
4. Routing del Contesto: Il contesto aggiornato dell'insieme viene instradato nuovamente ai singoli nodi. Gli autori propongono due varianti per questo passaggio:
   • GST-bc (Broadcast): Il token di grafo aggiornato $\tilde{z}_g$ viene trasmesso direttamente a tutti i nodi di quel grafo.
   • GST-ca (Cross-Attention): Ogni nodo presta attenzione all'insieme dei token di grafo aggiornati tramite un meccanismo di query.
5. Fusione Gated (Gated Fusion): Le caratteristiche originali del nodo, il contesto locale e il contesto dell'insieme vengono fusi utilizzando un meccanismo di gating asimmetrico appreso. Un MLP superficiale genera i gate che determinano quanta parte del contesto dell'insieme ($v_i$) rispetto al contesto locale ($u_i$) un nodo debba assorbire. Ciò consente ai nodi di ignorare il contesto dell'insieme se il loro grafo locale contiene già informazioni sufficienti.

Il modello completo consiste in uno stack di $L$ tali blocchi seguiti da una testa MLP per l'output specifico del compito.

Contributi Chiave

1. Innovazione Architetturale: Il documento propone una nuova architettura che fonde l'estrazione delle caratteristiche e la contestualizzazione dell'insieme ad ogni livello, contrastando con la standard pipeline "codifica-e-aggrega".
2. Meccanismo di Gating: L'introduzione di un meccanismo di gating asimmetrico consente una miscelazione adattiva di informazioni locali e globali, rispettando il bias induttivo per cui il grafo genitore di un nodo è sempre primario, mentre l'influenza dell'insieme è dipendente dal contesto.
3. Valutazione Esaustiva: Gli autori valutano GST su:
   • Benchmark Sintetici: Progettati per isolare il ragionamento strutturale condizionato dall'insieme (ad esempio, compiti RING-K, BALL-K, DIST e di Estrapolazione).
   • Applicazioni Chimiche: Previsione della resa di reazione (Buchwald-Hartwig) e identificazione del centro di reazione (USPTO-15K).
   • Classificazione di Immagini: Classificazione di insiemi CIFAR-10 (utilizzando immagini codificate come grafi).
4. Ablation e Analisi: Attraverso confronti con parametri accoppiati e ablazioni architettoniche, gli autori dimostrano che i guadagni di prestazione derivano dall'architettura intercalata piuttosto che da un aumento della capacità del modello.

Risultati

A parità di budget di parametri, GST supera costantemente i baseline (GCN + DeepSets e GCN + Set Transformer) in tutti i contesti valutati:

• Benchmark Sintetici: Le varianti di GST superano significativamente i baseline, in particolare man mano che la complessità del ragionamento condizionato dall'insieme aumenta (ad esempio, $K$ maggiore nei compiti RING-K). Il divario di prestazioni si amplia con l'aumentare degli spostamenti della dimensione dell'insieme, indicando che GST cattura le operazioni dell'insieme piuttosto che semplici statistiche del set di addestramento.
• Previsione della Resa di Reazione: GST raggiunge punteggi $R^2$ superiori, con i guadagni più significativi osservati in regimi con pochi dati e rumorosi (ad esempio, $R^2 = 0.53$ per GST-bc rispetto a $0.18$ per GCN-ST con il 2.5% di dati di addestramento).
• Previsione del Centro di Reazione: Sul benchmark USPTO-15K, GST-bc ottiene un vantaggio di 5.7 punti AP rispetto al miglior baseline a budget naturali di parametri. Il vantaggio persiste anche quando GST viene ridotto per corrispondere al numero di parametri del baseline o quando i baseline vengono espansi per corrispondere alla capacità di GST.
• Classificazione di Immagini: Sulla classificazione di insiemi CIFAR-10, le varianti di GST superano costantemente i baseline basati su GCN sia nelle metriche ROC-AUC che PR-AUC, con un divario che rimane significativo all'aumentare della dimensione dell'insieme.

Costo Computazionale: Gli autori riconoscono un compromesso; GST incorre in un costo di tempo di esecuzione (wall-clock) più elevato per ogni passo di addestramento (circa 4–6 volte quello dei baseline) a causa dell'attenzione a livello di insieme intercalata ad ogni livello. Tuttavia, la variante broadcast (GST-bc) offre un'alternativa più efficiente con una perdita di accuratezza trascurabile rispetto alla variante di cross-attention.

Significato e Rivendicazioni

Il documento sostiene che il significato primario di GST risieda nella sua capacità di rimuovere il collo di bottiglia tra la codifica del grafo e il ragionamento sull'insieme. Intercalando le informazioni locali e quelle dell'insieme, l'architettura permette agli elementi all'interno di un insieme di contestualizzarsi dinamicamente durante la fase di estrazione delle caratteristiche, invece di attendere una fase finale di pooling.

Gli autori sostengono che questo spostamento architettonico sia particolarmente cruciale per domini come la chimica, dove il comportamento di una molecola (o di un atomo all'interno di essa) è intrinsecamente dipendente dalla presenza e dalle proprietà di altre molecole nella miscela di reazione. I risultati suggeriscono che il "bias induttivo dell'intercalazione" sia un vantaggio fondamentale per compiti che richiedono un ragionamento strutturale condizionato dall'insieme, offrendo un'alternativa robusta alle pipeline sequenziali senza fare affidamento sulla pura capacità del modello per compensare i limiti architettonici.