Riepilogo Tecnico: GraViti – Autoencoder Variazionali a Livello di Grafo con Invarianza permutazionale Rilassata

Enunciato del Problema

L'autoencoding a livello di grafo mira a mappare interi grafi in vettori latenti compatti per facilitare compiti come l'interpolazione, la ricerca guidata dalle proprietà e la modellazione generativa. Sebbene gli Autoencoder Variazionali (VAE) a livello di nodo siano ben consolidati, essi soffrono di una limitazione fondamentale: lo spazio latente è definito a livello di nodo e dipende da un ordinamento arbitrario dei nodi. Di conseguenza, due grafi isomorfi possono mapparsi su insiemi di vettori latenti completamente diversi, impedendo la definizione di distanze significative tra grafi o l'acquisizione di varietà a livello di grafo.

Gli approcci esistenti a livello di grafo affrontano un compromesso. Metodi come PIGVAE e GRALE tentano di imporre una rigorosa invarianza permutazionale, affidandosi spesso a procedure di corrispondenza computazionalmente costose (ad esempio, matching di grafi, algoritmi di Sinkhorn) per allineare i grafi di input e ricostruiti. Queste fasi di allineamento introducono un significativo sovraccarico computazionale (tipicamente complessità cubica) e possono limitare la scalabilità su grandi dataset. Inoltre, molti decoder a livello di grafo si basano su strategie di padding che impongono una dimensione massima fissa del grafo o introducono classi di "nodi finti", limitando la generalizzazione a grafi più grandi e introducendo bias nei compiti di classificazione.

Metodologia: L'Architettura GraViti

GraViti è un VAE a livello di grafo basato su transformer progettato per codificare interi grafi in singoli vettori latenti, rilassando il requisito rigoroso di una ricostruzione invariante per permutazione quando esiste un ordinamento canonico coerente dei nodi.

1. Rappresentazione del Grafo

Il modello opera su grafi con classi discrete di nodi e archi (ad esempio, tipi di atomi, tipi di legami).

• Input: Una matrice di classi di nodi $X$, un tensore di classi di archi $E$ e un vettore di caratteristiche globali $y$ che codifica statistiche a livello di grafo (conteggi di tipi di nodi/archi e dimensione del grafo).
• Codifiche Posizionali: Per iniettare informazioni strutturali, il modello utilizza autovettori del Laplaciano e caratteristiche di cammini casuali come codifiche posizionali ($P$).
• Conoscenza del Dominio: Per domini specifici come le molecole, vengono incorporate caratteristiche aggiuntive specifiche per nodo ($K$) come massa atomica o elettronegatività.
• Batching: I grafi vengono riempiti (padded) fino alla dimensione del grafo più grande in un batch, ma il modello è progettato per ignorare completamente questi elementi riempiti, evitando la necessità di classi di nodi finti.

2. Encoder

L'encoder mappa la tripletta di input $(X, E, y)$ nei parametri della distribuzione latente $(z_\mu, z_\sigma)$.

• Embedding: Le classi di nodi, le classi di archi e le caratteristiche globali vengono proiettate in embedding.
• Stack di Transformer per Grafi: Gli embedding vengono elaborati attraverso uno stack di livelli di transformer per grafi (basati sulle architetture DiGress e CatFlow). Ciò consente aggiornamenti simultanei di caratteristiche a livello di nodo, arco e grafo, garantendo l'equivalenza permutazionale durante la fase di codifica.
• Readout: A differenza del pooling somma/media tradizionale, GraViti impiega un readout basato su transformer sequenziale. Due decoder transformer sequenziali paralleli elaborano rispettivamente le rappresentazioni di nodi e archi, utilizzando il vettore di caratteristiche globali $y_{rep}$ come target per iniettare il contesto globale. Ciò garantisce che l'output sia indipendente dall'ordinamento dei nodi nella fase di readout.

3. Predizione della Dimensione del Grafo

Un'innovazione cruciale è la disentanglement esplicita della dimensione del grafo dal resto della rappresentazione latente.

• Il modello forza le prime $d_s$ dimensioni dello spazio latente a codificare il log-dimensione del grafo.
• Un MLP predice il log-dimensione $\hat{s}$ da queste dimensioni.
• Durante la decodifica, il log-dimensione reale $s$ viene ricodificato e reiniettato nella rappresentazione latente per guidare la generazione del numero corretto di nodi, assicurando che il decoder conosca la dimensione target del grafo senza fare affidamento su maschere.

4. Decoder

Il decoder ricostruisce il grafo dal vettore latente $z$.

• Broadcasting: Il vettore latente viene duplicato $N$ volte (dove $N$ è la dimensione predetta o reale) e arricchito con codifiche posizionali sinusoidali per rompere la simmetria.
• Transformer Sequenziale: Un encoder transformer sequenziale elabora questa sequenza, dividendola in sequenze orientate ai nodi e orientate agli archi.
• Raffinamento: Un secondo stack di transformer per grafi raffina le rappresentazioni di nodi e archi.
• Teste di Predizione: MLP finali classificano i tipi di nodi e i tipi di archi. Il modello predice direttamente la dimensione del grafo, eliminando la necessità di una classe finta "nessun arco".

5. Obiettivo di Addestramento

Il modello è addestrato utilizzando un obiettivo VAE standard:

• Perdita di Ricostruzione: Una combinazione di Focal Loss per la classificazione di nodi e archi (per gestire lo squilibrio delle classi) e una perdita L1 smooth per la predizione della dimensione del grafo.
• Perdita Variazionale: Un termine di divergenza KL regolarizza lo spazio latente verso un prior normale standard.
• Invarianza Permutazionale Rilassata: Crucialmente, la perdita di ricostruzione viene applicata in base all'ordinamento dei nodi di input. Gli autori sostengono che in domini con un ordinamento canonico affidabile (ad esempio, molecole, reti bayesiane), imporre una rigorosa invarianza permutazionale tramite matching è dannoso. Affidandosi a un ordinamento coerente, il modello raggiunge una complessità quadratica e una maggiore fedeltà di ricostruzione.

Contributi Chiave

1. Architettura Scalabile: GraViti raggiunge un'accuratezza di ricostruzione all'avanguardia su grandi dataset di grafi (PubChem16, PubChem32) evitando costose procedure di matching invariate per permutazione.
2. Spazio Latente Strutturato: Lo spazio latente appreso risulta ben strutturato, supportando un'interpolazione fluida, una modifica controllata e un'ottimizzazione guidata dalle proprietà. Su benchmark molecolari, il modello recupera principi chimici consolidati (ad esempio, la relazione tra lunghezza della catena alchilica e idrofobicità) direttamente dallo spazio latente.
3. Invarianza Permutazionale Rilassata: Il documento dimostra empiricamente che l'invarianza permutazionale rigorosa non è sempre necessaria. Quando è disponibile un ordinamento canonico coerente dei nodi, rilassare questo vincolo migliora significativamente la qualità della ricostruzione e la scalabilità rispetto ai metodi che impongono l'invarianza attraverso matching complessi.

Risultati Sperimentali

Gli autori hanno valutato GraViti su diversi benchmark:

• Ricostruzione: Su dataset molecolari (QM9, PubChem16, PubChem32), GraViti supera i baseline (GRALE, PIGVAE) fino al 25% in termini di accuratezza di ricostruzione, in particolare su grafi più grandi dove i metodi basati sul matching faticano.
• Studi di Ablazione: La rimozione dell'encoder transformer o del readout basato su transformer degrada significativamente le prestazioni, confermando l'importanza di questi componenti.
• Denoising: Il modello ripristina efficacemente strutture chimicamente valide da input rumorosi, con la variante VAE che mostra una robustezza superiore rispetto alla variante AE.
• Ottimizzazione delle Proprietà: Utilizzando CMA-ES nello spazio latente, il modello massimizza con successo il logP (lipofilia), generando molecole con catene di carbonio progressivamente più lunghe o specifici gruppi idrofobici, allineandosi alle tendenze chimiche note.
• Interpolazione: L'interpolazione lineare nello spazio latente produce transizioni fluide tra strutture molecolari.
• Compiti a valle: I vettori latenti addestrati su PubChem predicono con alta accuratezza le proprietà molecolari (momento di dipolo, gap HOMO-LUMO) e i punteggi del Criterio di Informazione Bayesiano (BIC) per le reti bayesiane.
• Generazione: Sebbene la validità della generazione sia inferiore su PubChem (circa 50-70% a seconda della rappresentazione aromatica) rispetto a QM9, il modello produce molecole diverse e nuove. Gli autori attribuiscono la validità inferiore su dataset più grandi alla compattezza globale dello spazio latente piuttosto che a fallimenti strutturali locali.

Significato e Affermazioni

Il documento afferma che GraViti fornisce una base pratica per l'apprendimento di rappresentazioni a livello di grafo. Il suo significato principale risiede nel dimostrare che l'invarianza permutazionale rigorosa non è un prerequisito per un autoencoding a livello di grafo efficace quando esiste un ordinamento stabile dei nodi. Rilassando questo vincolo, GraViti evita i colli di bottiglia computazionali del matching di grafi, consentendo un addestramento scalabile su dataset massicci mantenendo al contempo un'alta qualità di ricostruzione.

Gli autori posizionano GraViti come un'alternativa leggera a pipeline di generazione complesse, offrendo un meccanismo di decodifica in un singolo passaggio che bilancia le prestazioni generative con la qualità pratica dei campioni. Il lavoro suggerisce che per domini come i grafi molecolari e le reti bayesiane, dove gli ordinamenti canonici sono intrinseci o possono essere costruiti, i modelli a livello di grafo possono raggiungere prestazioni superiori senza il sovraccarico di perdite basate sull'allineamento. Si propone un lavoro futuro per integrare vincoli consapevoli del dominio per migliorare ulteriormente la validità della generazione ed esplorare il trasferimento ad altri domini di grafi come circuiti e proteine.