Riepilogo Tecnico: Diffusione Latente Principiata per Grafi tramite Autoencoder Laplaciani

Enunciato del Problema

I modelli di diffusione per grafi hanno raggiunto prestazioni all'avanguardia in compiti generativi in domini quali la progettazione molecolare, l'ottimizzazione combinatoria e la ricerca di architetture neurali. Tuttavia, gli approcci esistenti operano tipicamente direttamente nello spazio dei grafi, sostenendo una complessità computazionale quadratica ($O(n^2)$) rispetto al numero di nodi. Questo collo di bottiglia nasce dal fatto che questi modelli devono apprendere sia le rappresentazioni dei nodi che quelle degli archi, spesso sprecando capacità nel modellare l'assenza di archi in grafi sparsi. Inoltre, i tentativi di adattare modelli autoregressivi scalabili ai grafi spesso rompono l'equivarianza permutazionale linearizzando le strutture dei grafi.

Un'alternativa promettente è la diffusione latente, che comprime i dati in uno spazio a bassa dimensionalità prima di applicare il processo generativo (come osservato nella generazione di immagini con VAE). Tuttavia, applicare questo concetto ai grafi presenta una sfida unica: a differenza delle immagini, dove errori di decodifica minori sono visivamente trascurabili, la generazione di grafi richiede una ricostruzione quasi senza perdite. Un singolo errore nella ricostruzione di una matrice di adiacenza può invalidare l'intero campione (ad esempio, rompendo legami chimici o violando vincoli strutturali). I modelli esistenti di diffusione latente per grafi spesso non soddisfano questo requisito di fedeltà o sacrificano l'efficienza utilizzando architetture ad hoc che dipendono ancora da token di archi quadratici.

Metodologia: LG-Flow

Gli autori propongono LG-Flow, un framework di diffusione latente per grafi progettato per superare questi ostacoli mantenendo un'elevata fedeltà di ricostruzione. Il framework è composto da due componenti principali: un autoencoder principiato e un modello generativo di flusso latente.

1. Autoencoder Variazionale per Grafi Laplaciano (LG-VAE)

L'innovazione centrale è un autoencoder equivariante per permutazione che mappa i grafi in uno spazio latente in cui la dimensionalità scala linearmente con il numero di nodi ($n \times d$), piuttosto che quadraticamente.

• Base Teorica: L'architettura è fondata su codici di posizione identificativi dell'adiacenza. Gli autori sfruttano la proprietà secondo cui, se gli embedding dei nodi sono sufficientemente "identificativi dell'adiacenza", la matrice di adiacenza può essere ricostruita a partire da questi embedding.
• Codificatore: Il codificatore utilizza Codici di Posizione Laplaciani (LPE) derivati dalla decomposizione spettrale del Laplaciano del grafo ($L = U\Lambda U^\top$). Incorpora le caratteristiche dei nodi e degli archi in una distribuzione a posteriori gaussiana. Per gestire i Grafi Aciclici Diretti (DAG), il framework si estende al Laplaciano Magnetico ($L_q$), una matrice hermitiana che preserva la direzionalità degli archi, risultando in un LPE Magnetico (mLPE).
• Decodificatore: Il decodificatore è progettato per una ricostruzione ad alta fedeltà.
  • Etichette dei Nodi: Decodificate direttamente tramite un layer lineare e softmax.
  • Matrice di Adiacenza: Ricostruita tramite un layer bilineare applicato agli embedding latenti ($Z W_Q (Z W_K)^\top$). La rilevazione degli archi è trattata come un problema di classificazione binaria utilizzando un'architettura DeepSet per garantire l'equivarianza permutazionale e gestire dimensioni variabili dei grafi.
  • Etichette degli Archi: Decodificate utilizzando un layer bilineare multi-testa.
• Garanzia di Ricostruzione: Il design assicura che le rappresentazioni latenti contengano informazioni strutturali sufficienti per recuperare la matrice di adiacenza, affrontando la necessità critica di una ricostruzione quasi senza perdite nella generazione di grafi.

2. Flusso Latente di Adattamento (LG-Flow)

Una volta addestrato il LG-VAE, i suoi parametri vengono congelati. Successivamente, un modello generativo viene addestrato sullo spazio latente $Z \in \mathbb{R}^{n \times d}$.

• Architettura: Gli autori impiegano un Trasformatore di Diffusione (DiT) con obiettivi di Flusso di Adattamento (FM).
• Processo: Il modello impara a prevedere il campo di velocità che trasporta il rumore verso la distribuzione latente.
• Efficienza: Operando nello spazio latente, il DiT evita la complessità quadratica di prevedere direttamente l'intera matrice di adiacenza. L'unico passaggio quadratico è la decodifica finale, gestita da un decodificatore leggero. Ciò consente l'addestramento di backbones generativi sostanzialmente più grandi.

Contributi Chiave

1. LG-VAE: Un autoencoder principiato ed equivariante per permutazione che raggiunge una ricostruzione quasi senza perdite sia per grafi non diretti che per DAG. È il primo a fornire garanzie teoriche (tramite proprietà identificative dell'adiacenza) che la matrice di adiacenza è recuperabile dallo spazio latente.
2. Scalabilità Lineare: Il framework mappa i nodi in embedding a dimensionalità fissa, assicurando che la dimensione della rappresentazione latente scala linearmente con il numero di nodi, rimuovendo efficacemente il collo di bottiglia quadratico della diffusione nello spazio dei grafi.
3. Framework Unificato: L'approccio unifica la generazione di grafi non diretti e DAG all'interno di un singolo framework modulare, consentendo l'uso di architetture generiche di denoising (DiT) e tecniche derivanti dalla generazione di immagini.
4. Prestazioni: Il metodo raggiunge prestazioni generative competitive fornendo al contempo significativi accelerazioni (fino a 1000$\times$) rispetto ai modelli di diffusione per grafi all'avanguardia.

Risultati Sperimentali

Gli autori hanno valutato LG-Flow su sei dataset: tre sintetici (Extended Planar, Extended Tree, Ego), due molecolari (MOSES, GuacaMol) e due dataset di DAG (TPU Tile, Price).

• Fedeltà di Ricostruzione (Q1): Su benchmark strutturali (Extended Planar e Extended Tree), il LG-VAE ha raggiunto una ricostruzione quasi senza perdite (Accuratezza del Campione di 0,9961 e 0,9844 rispettivamente), superando significativamente le linee di base come DGVAE e GLAD, che non sono riuscite a ricostruire i campioni esattamente.
• Prestazioni Generative (Q2):
  • Grafici Sintetici: LG-Flow ha ottenuto risultati competitivi sulle metriche di distribuzione (Grado, Orbita, Clustering, Spettro) e su Validità, Unicità, Novità (V.U.N.).
  • Molecole: Su GuacaMol e MOSES, il metodo ha raggiunto prestazioni all'avanguardia o competitive in termini di validità e corrispondenza della distribuzione chimica (FCD), superando spesso i modelli di diffusione discreta come DiGress e DisCo.
  • DAG: Sul dataset TPU Tile, LG-Flow ha superato LayerDag autoregressivo e Directo basato su diffusione sulle metriche di distribuzione, mantenendo un'elevata validità.
• Efficienza (Q3):
  • LG-Flow ha dimostrato accelerazioni massive, con inferenze da 10$\times$ a 1000$\times$ più veloci rispetto a linee di base come DeFoG e Cometh.
  • L'utilizzo della memoria è scalato linearmente con la dimensione del batch, permettendo all'intero batch di test di 151 campioni di adattarsi su una singola GPU, mentre DeFoG esauriva la memoria con dimensioni del batch basse come 32.

Significato e Affermazioni

Il documento afferma che LG-Flow dimostra che spostare la diffusione dei grafi in uno spazio latente di nodi ad alta fedeltà può migliorare sostanzialmente l'efficienza dell'inferenza, in particolare per grafi piccoli e medi. Il framework sposta il focus dalla progettazione di modelli di diffusione complessi e specifici per i grafi allo sviluppo di autoencoder per grafi potenti e più economici da addestrare.

Unificando la generazione di grafi non diretti e diretti e permettendo il trasferimento di architetture potenti come il Trasformatore di Diffusione dalla generazione di immagini, gli autori suggeriscono una strada verso la scalabilità dei modelli generativi a grafi molto più grandi che in precedenza erano intrattabili a causa della complessità quadratica. Il lavoro evidenzia che, con un autoencoder principiato che garantisce la fedeltà di ricostruzione, i benefici della diffusione latente (efficienza, modularità e capacità) possono essere realizzati con successo nel dominio dei grafi.