Toward Better Geometric Representations for Molecule Generative Models

Technische Samenvatting: LENSEs – Latente Versterking voor Niet-Gladde Structurele Encoderings

Probleemstelling
Geometrisch representatie-geconditioneerde generatie (GeoRCG) heeft zich ontwikkeld als een efficiënt paradigma voor 3D-moleculengeneratie, waarbij de modellering van moleculaire representaties wordt ontkoppeld van de generatie van 3D-structuren. In dit kader definieert een vooraf getrainde moleculaire encoder (bijv. UniMol, Frad) een latente ruimte, een Representatie Diffusiemodel (RDM) steekt een molecuulniveau-vector uit deze ruimte, en een generator reconstrueert de 3D-structuur geconditioneerd op deze vector.

De effectiviteit van GeoRCG wordt echter fundamenteel beperkt door de kwaliteit van de latente ruimte die door vooraf getrainde encoders wordt geboden. De auteurs identificeren twee kritieke knelpunten:

1. Niet-gladde Representatievariëteiten: Vooraf getrainde encoders, die primair worden geoptimaliseerd voor eigendomsvoorspelling via coördinatenontruising, produceren representaties met hoge Lipschitz-constanten. Kleine verstoringen in moleculaire conformatie leiden tot disproportioneel grote verschuivingen in de latente ruimte. Deze niet-gladheid maakt de verdeling moeilijk te modelleren voor de RDM, wat vaak compensatiemechanismen vereist zoals handmatig afgestemde ruisinjectie.
2. Onderbenutte Multi-niveau Semantiek: Vooraf getrainde encoders vertonen hiërarchische kenmerken waarbij ondiepe lagen lokale structurele motieven (bijv. functionele groepen) vastleggen en diepere lagen globale semantiek coderen. Bestaande GeoRCG-pijplijnen gebruiken doorgaans alleen de representatie van de laatste laag, waardoor de rijke semantische informatie die in eerdere lagen wordt bewaard, wordt weggegooid. Bijgevolg moet de generator deze semantiek vanaf nul opnieuw leren, wat de optimalisatie moeilijkheid verhoogt.

Methodologie: LENSEs
Om deze problemen aan te pakken, stellen de auteurs LENSEs (Latente Versterking voor Niet-Gladde Structurele Encoderings) voor, een kader dat actief de conditionerende latente ruimte verfijnt tijdens de generatieve training. In plaats van passief te vertrouwen op een bevroren encoder, introduceert LENSEs drie complementaire mechanismen:

1. Representatiekop met Leerbaar Laagpoolsen:
   In plaats van een enkele deterministische output van de encoder te gebruiken, aggregeert LENSEs kenmerken uit meerdere lagen. Het maakt gebruik van een leerbaar laagpoolsmechanisme om representaties uit verschillende encoderdieptes ($H^{(l)}$) te wegen. Deze gepoolde kenmerken worden via projectiekoppen gemapt naar de parameters van een diagonale Gaussische posterior ($\mu, \sigma$). De latente code $z$ wordt bemonsterd met behulp van de herparameteriseringstruc. Deze kop wordt gezamenlijk met de generator getraind om representaties te extraheren die zowel semantisch rijk als generatievriendelijk zijn, geregulariseerd door een KL-divergentie-verlies om een compacte latente ruimte te waarborgen.

2. Molecuul Perceptueel Verlies ($L_{PERC}$):
   Geïnspireerd door perceptuele verliezen in computer vision, aligneert dit verlies de hoog-niveau semantische representatie van het ontroeste (gegenereerde) molecuul met die van het schone grondwaarheid-molecuul binnen de kenmerkrimte van de bevroren encoder. Door zowel het gegenereerde als het schone molecuul door de encoder te leiden en de afstand tussen hun gepoolde graf-niveau representaties te minimaliseren, wordt de generator aangemoedigd betekenisvolle structurele patronen (bijv. functionele groepen) te behouden in plaats van alleen lage-niveau coördinaatfouten te minimaliseren. Een stop-gradient-operatie wordt toegepast op de grondwaarheid-representatie.

3. Representatie-Alignement (REPA) Verlies ($L_{REPA}$):
   Gemotiveerd door de observatie dat verborgen toestanden van de generator van nature aligneren met encoder-kenmerken, aligneert dit verlies expliciet de knooppuntniveau-kenmerken van de generator met die van de vooraf getrainde encoder. Een leerbare projector mapt generatorkenmerken naar de kenmerkrimte van de encoder, en de cosinus-ähnelijkheid wordt gemaximaliseerd tussen overeenkomstige knooppunten. Dit vermindert de semantische kloof tussen vooraf trainen en generatie, versnelt de training en verbetert de acquisitie van chemisch betekenisvolle kennis.

Het totale trainingsdoel combineert het basisgeneratieverlies ($L_{gen}$) met het KL-verlies, perceptuele verlies en REPA-verlies.

Belangrijkste Resultaten
De auteurs evalueren LENSEs op de GEOM-DRUG- en QM9-datasets, waarbij ze het vergelijken met state-of-the-art baselines, waaronder GeoRCG, CanonFlow en SemlaFlow.

• Generatiekwaliteit: Op de uitdagende GEOM-DRUG-dataset bereikt LENSEs state-of-the-art prestaties. Het behaalt 97,28% validiteit en 98,51% molecuulstabiliteit, en overtreft de vorige beste GeoRCG (95,3% validiteit, 97,20% stabiliteit) en CanonFlow (95,9% validiteit). Het bereikt ook de laagste energie- en spanningsmetrieken, wat wijst op superieure fysieke plausibiliteit.
• Representatiegladheid: Analyse van de Lipschitz-constante onthult dat de verfijnde latente ruimte aanzienlijk gladder is dan de ruwe encoder-ruimte. Op GEOM-DRUG wordt de empirische Lipschitz-constante met 4,6× verlaagd, wat bevestigt dat de representatievariëteit beter geschikt is voor diffusiemodellering.
• Representatierijkdom: De effectieve rang van de representatie stijgt van 65,32 (ruwe encoder) naar 88,42 (LENSEs), wat wijst op een vermindering van dimensionale instorting en een rijkere, meer gestructureerde latente ruimte.
• Downstream Nut: Wanneer gebruikt voor downstream regressietaken op QM9-eigenschappen met een bevroren encoder en representatiekop, presteert LENSEs beter dan de ruwe encoder voor alle eigenschappen, wat aantoont dat het generatieve trainingsproces zelf fungeert als een effectief voortrainingsparadigma voor het verfijnen van moleculaire representaties.
• Bemonsteringsefficiëntie: LENSEs toont verbeterde bemonsteringsefficiëntie, waarbij hoge validiteit en stabiliteit worden bereikt met minder generatiestappen in vergelijking met de baseline.

Betekenis en Claims
Het artikel stelt dat hoogwaardige 3D-moleculengeneratie niet alleen afhangt van de generatorarchitectuur, maar kritiek van de kwaliteit van de conditionerende representatie. LENSEs toont aan dat de latente ruimte van vooraf getrainde encoders systematisch kan worden verbeterd om gladder, semantisch informatiever en beter afgestemd te zijn op generatieve doelen.

De auteurs stellen dat hun werk een bidirectionele relatie tussen generatiekwaliteit en representatiekwaliteit vestigt. Specifiek suggereren ze dat generatieve training met alignementdoelen kan dienen als een potentieel voortrainingsparadigma voor moleculaire encoders, waarbij representaties worden verfijnd om voorspellender en nuttiger te zijn voor downstream-taken. Het werk benadrukt dat nuttige moleculaire representaties niet van nature generatievriendelijk zijn en expliciete verfijning vereisen om hun volledige potentieel in generatieve modellering te ontsluiten.