Technische Samenvatting: Principiële Latente Diffusie voor Grafen via Laplace-automatische Encoders

Probleemstelling

Grafdiffusiemodellen hebben state-of-the-art prestaties behaald in generatieve taken op gebieden zoals moleculair ontwerp, combinatorische optimalisatie en zoektocht naar neurale architecturen. Echter, bestaande benaderingen werken doorgaans direct in de grafruimte, wat leidt tot kwadratische computationele complexiteit ($O(n^2)$) met betrekking tot het aantal knopen. Deze bottleneck ontstaat omdat deze modellen zowel knoop- als randrepresentaties moeten leren, vaak met een verspilling van capaciteit bij het modelleren van de afwezigheid van randen in schaarse grafen. Bovendien breken pogingen om schaalbare autoregressieve modellen aan grafen aan te passen vaak de permutatie-equivariantie door grafstructuren te lineariseren.

Een veelbelovend alternatief is latente diffusie, waarbij data eerst wordt gecomprimeerd naar een laag-dimensionale ruimte voordat het generatieve proces wordt toegepast (zoals gezien bij beeldgeneratie met VAE's). Het toepassen hiervan op grafen vormt echter een unieke uitdaging: in tegenstelling tot afbeeldingen, waar kleine decoderingsfouten visueel verwaarloosbaar zijn, vereist grafgeneratie een bijna-verliesvrije reconstructie. Een enkele fout bij het reconstrueren van een burenmatrix kan het hele voorbeeld ongeldig maken (bijvoorbeeld door chemische bindingen te verbreken of structurele beperkingen te schenden). Bestaande latente grafdiffusiemodellen voldoen vaak niet aan deze vereiste van fideliteit of offeren efficiëntie op door het gebruik van ad-hoc architecturen die nog steeds afhankelijk zijn van kwadratische randtokens.

Methodologie: LG-Flow

De auteurs stellen LG-Flow voor, een latent grafdiffusiekader dat is ontworpen om deze obstakels te overwinnen terwijl een hoge reconstructiefideliteit wordt behouden. Het kader bestaat uit twee hoofdcomponenten: een principiële automatische encoder en een generatief model voor latente stroommatching.

1. Laplace Graf Variational Autoencoder (LG-VAE)

De kerninnovatie is een permutatie-equivariante automatische encoder die grafen afbeeldt naar een latente ruimte waar de dimensionaliteit lineair schaalt met het aantal knopen ($n \times d$), in plaats van kwadratisch.

• Theoretische Basis: De architectuur is gebaseerd op positional encodings die buren identificeren. De auteurs maken gebruik van het gegeven dat, als knoopembeddings voldoende "buren-identificerend" zijn, de burenmatrix uit deze embeddings kan worden gereconstrueerd.
• Encoder: De encoder maakt gebruik van Laplace Positional Encodings (LPE) afgeleid van de eigendecompositie van de graf-Laplace-operator ($L = U\Lambda U^\top$). Het integreert knoop- en randkenmerken in een Gaussische posterieure verdeling. Om Gerichte Acyclische Grafen (DAG's) te verwerken, wordt het kader uitgebreid tot de magnetische Laplace-operator ($L_q$), een Hermitiaanse matrix die de randrichting behoudt, wat resulteert in een magnetische LPE (mLPE).
• Decoder: De decoder is ontworpen voor reconstructie met hoge fideliteit.
  • Knooplabels: Direct gedecodeerd via een lineaire laag en softmax.
  • Burenmatrix: Gereconstrueerd via een bilineaire laag toegepast op de latente embeddings ($Z W_Q (Z W_K)^\top$). Het detecteren van randen wordt behandeld als een binaire classificatieprobleem met behulp van een DeepSet-architectuur om permutatie-equivariantie te waarborgen en variabele grafgroottes te hanteren.
  • Randlabels: Gedecodeerd met behulp van een multi-headed bilineaire laag.
• Reconstructiegarantie: Het ontwerp zorgt ervoor dat de latente representaties voldoende structurele informatie bevatten om de burenmatrix te herstellen, waarmee wordt ingegaan op de kritieke behoefte aan bijna-verliesvrije reconstructie in grafgeneratie.

2. Latente Stroommatching (LG-Flow)

Zodra de LG-VAE is getraind, worden zijn parameters bevroren. Vervolgens wordt een generatief model getraind op de latente ruimte $Z \in \mathbb{R}^{n \times d}$.

• Architectuur: De auteurs maken gebruik van een Diffusion Transformer (DiT) met Flow Matching (FM)-doelstellingen.
• Proces: Het model leert het snelheidsveld te voorspellen dat ruis transporteert naar de latente verdeling.
• Efficiëntie: Door te opereren in de latente ruimte, vermijdt de DiT de kwadratische complexiteit van het direct voorspellen van de volledige burenmatrix. De enige kwadratische stap is de uiteindelijke decoding, die wordt afgehandeld door een lichtgewicht decoder. Dit maakt het mogelijk om aanzienlijk grotere generatieve backbones te trainen.

Belangrijkste Bijdragen

1. LG-VAE: Een principiële, permutatie-equivariante automatische encoder die bijna-verliesvrije reconstructie bereikt voor zowel ongerichte grafen als DAG's. Het is de eerste die theoretische garanties biedt (via eigenschappen van buren-identificatie) dat de burenmatrix uit de latente ruimte herstelbaar is.
2. Lineaire Schaalbaarheid: Het kader beeldt knopen af op embeddings met een vaste dimensionaliteit, waardoor de grootte van de latente representatie lineair schaalt met het aantal knopen, wat effectief de kwadratische bottleneck van diffusie in grafruimte verwijdert.
3. Unificerend Kader: De benadering verenigt de generatie van ongerichte grafen en DAG's binnen één modulair kader, waardoor het gebruik mogelijk wordt van generieke denoising-architecturen (DiT) en technieken uit beeldgeneratie.
4. Prestaties: De methode bereikt concurrerende generatieve prestaties terwijl het aanzienlijke snelheidswinsten levert (tot 1000$\times$) in vergelijking met state-of-the-art grafdiffusiemodellen.

Experimentele Resultaten

De auteurs hebben LG-Flow geëvalueerd op zes datasets: drie synthetische (Extended Planar, Extended Tree, Ego), twee moleculaire (MOSES, GuacaMol) en twee DAG-datasets (TPU Tile, Price).

• Reconstructiefideliteit (Q1): Op structurele benchmarks (Extended Planar en Extended Tree) bereikte LG-VAE bijna-verliesvrije reconstructie (Stichproefnauwkeurigheid van respectievelijk 0,9961 en 0,9844), wat aanzienlijk beter was dan baselines zoals DGVAE en GLAD, die niet in staat waren om voorbeelden exact te reconstrueren.
• Generatieve Prestaties (Q2):
  • Synthetische Grafen: LG-Flow behaalde concurrerende resultaten op verdelingsmetrieken (Graad, Orbit, Clustering, Spectrum) en Validiteit, Uniekheid, Nieuwigheid (V.U.N.).
  • Moleculen: Op GuacaMol en MOSES behaalde de methode state-of-the-art of concurrerende prestaties op het gebied van validiteit en chemische verdelingsmatching (FCD), vaak presterend boven discrete diffusiemodellen zoals DiGress en DisCo.
  • DAG's: Op de TPU Tile-dataset overtrof LG-Flow de autoregressieve LayerDAG en de op diffusie gebaseerde Directo op verdelingsmetrieken, terwijl het een hoge validiteit behield.
• Efficiëntie (Q3):
  • LG-Flow toonde enorme snelheidswinsten, variërend van 10$\times$ tot 1000$\times$ snellere inferentie in vergelijking met baselines zoals DeFoG en Cometh.
  • Het geheugengebruik schaalde lineair met de batchgrootte, waardoor de volledige testbatch van 151 voorbeelden op één GPU paste, terwijl DeFoG al bij batchgroottes zo laag als 32 uit het geheugen liep.

Betekenis en Claims

Het artikel beweert dat LG-Flow aantoont dat het verplaatsen van grafdiffusie naar een latente ruimte met hoge fideliteit op knoopniveau de inferentie-efficiëntie aanzienlijk kan verbeteren, met name voor kleine tot middelgrote grafen. Het kader verschuift de focus van het ontwerpen van complexe, graf-specifieke diffusiemodellen naar het ontwikkelen van sterke, goedkopere te trainen grafautomatische encoders.

Door de generatie van ongerichte en gerichte grafen te verenigen en het overdragen van krachtige architecturen zoals de Diffusion Transformer vanuit beeldgeneratie mogelijk te maken, suggereren de auteurs een weg naar het schalen van generatieve modellen naar veel grotere grafen die eerder onbereikbaar waren vanwege kwadratische complexiteit. Het werk benadrukt dat met een principiële automatische encoder die reconstructiefideliteit waarborgt, de voordelen van latente diffusie (efficiëntie, modulariteit en capaciteit) succesvol kunnen worden gerealiseerd in het domein van grafen.