Toward Better Geometric Representations for Molecule Generative Models

Sintesi Tecnica: LENSEs – Miglioramento Latente per Codifiche Strutturali Non Lisce

Enunciato del Problema
La generazione condizionata dalla rappresentazione geometrica (GeoRCG) si è affermata come un paradigma efficiente per la generazione di molecole 3D, disaccoppiando la modellazione delle rappresentazioni molecolari dalla generazione delle strutture 3D. In questo quadro, un encoder molecolare preaddestrato (ad esempio, UniMol, Frad) definisce uno spazio latente, un Modello di Diffusione di Rappresentazione (RDM) campiona un vettore a livello di molecola da questo spazio, e un generatore ricostruisce la struttura 3D condizionata da tale vettore.

Tuttavia, l'efficacia di GeoRCG è fondamentalmente limitata dalla qualità dello spazio latente fornito dagli encoder preaddestrati. Gli autori identificano due colli di bottiglia critici:

1. Varietà di Rappresentazione Non Lisce: Gli encoder preaddestrati, ottimizzati principalmente per la previsione di proprietà tramite denoising delle coordinate, producono rappresentazioni con costanti di Lipschitz elevate. Piccole perturbazioni nella conformazione molecolare portano a spostamenti sproporzionatamente grandi nello spazio latente. Questa non liscitudine rende difficile per l'RDM modellare la distribuzione, richiedendo spesso meccanismi compensativi come l'iniezione manuale di rumore.
2. Semantica Multi-livello Sfruttata Inadeguatamente: Gli encoder preaddestrati esibiscono caratteristiche gerarchiche dove i livelli superficiali catturano motivi strutturali locali (ad esempio, gruppi funzionali) e i livelli più profondi codificano semantica globale. Le pipeline GeoRCG esistenti utilizzano tipicamente solo la rappresentazione dell'ultimo livello, scartando le ricche informazioni semantiche preservate nei livelli precedenti. Di conseguenza, il generatore deve reimparare queste semantica da zero, aumentando la difficoltà di ottimizzazione.

Metodologia: LENSEs
Per affrontare questi problemi, gli autori propongono LENSEs (Miglioramento Latente per Codifiche Strutturali Non Lisce), un quadro che affina attivamente lo spazio latente di condizionamento durante l'addestramento generativo. Piuttosto che affidarsi passivamente a un encoder congelato, LENSEs introduce tre meccanismi complementari:

1. Testa di Rappresentazione con Pooling di Livelli Apprendibile:
   Invece di utilizzare un singolo output deterministico dall'encoder, LENSEs aggrega caratteristiche multilivello. Impiega un meccanismo di pooling di livelli apprendibile per pesare le rappresentazioni da diverse profondità dell'encoder ($H^{(l)}$). Queste caratteristiche aggregate sono mappate sui parametri di una posteriora Gaussiana diagonale ($\mu, \sigma$) tramite teste di proiezione. Il codice latente $z$ è campionato utilizzando il trucco di riparametrizzazione. Questa testa è addestrata congiuntamente al generatore per estrarre rappresentazioni che siano sia semanticamente ricche che favorevoli alla generazione, regolarizzate da una perdita di divergenza KL per garantire uno spazio latente compatto.

2. Perdita Percettiva Molecolare ($L_{PERC}$):
   Ispirata alle perdite percettive nella visione artificiale, questa perdita allinea la rappresentazione semantica di alto livello della molecola denoizzata (generata) con quella della molecola pulita di verità fondamentale (ground-truth) nello spazio delle caratteristiche dell'encoder congelato. Facendo passare sia la molecola generata che quella pulita attraverso l'encoder e minimizzando la distanza tra le loro rappresentazioni aggregate a livello di grafo, si incoraggia il generatore a preservare modelli strutturali significativi (ad esempio, gruppi funzionali) piuttosto che semplicemente minimizzare errori di coordinate di basso livello. Un'operazione di stop-gradient viene applicata alla rappresentazione della verità fondamentale.

3. Perdita di Allineamento di Rappresentazione (REPA) ($L_{REPA}$):
   Motivata dall'osservazione che gli stati nascosti del generatore si allineano naturalmente con le caratteristiche dell'encoder, questa perdita allinea esplicitamente le caratteristiche a livello di nodo del generatore con quelle dell'encoder preaddestrato. Un proiettore apprendibile mappa le caratteristiche del generatore nello spazio delle caratteristiche dell'encoder, e la similarità del coseno è massimizzata tra i nodi corrispondenti. Questo riduce il divario semantico tra preaddestramento e generazione, accelerando l'addestramento e migliorando l'acquisizione di conoscenze chimicamente significative.

L'obiettivo totale di addestramento combina la perdita di generazione di base ($L_{gen}$) con la perdita KL, la perdita percettiva e la perdita REPA.

Risultati Chiave
Gli autori valutano LENSEs sui dataset GEOM-DRUG e QM9, confrontandolo con baseline all'avanguardia tra cui GeoRCG, CanonFlow e SemlaFlow.

• Qualità della Generazione: Sul dataset impegnativo GEOM-DRUG, LENSEs ottiene prestazioni all'avanguardia. Raggiunge una validità del 97,28% e una stabilità molecolare del 98,51%, superando il precedente miglior GeoRCG (95,3% di validità, 97,20% di stabilità) e CanonFlow (95,9% di validità). Ottiene anche le metriche di energia e deformazione più basse, indicando una plausibilità fisica superiore.
• Liscitudine della Rappresentazione: L'analisi della costante di Lipschitz rivela che lo spazio latente raffinato è significativamente più liscio dello spazio grezzo dell'encoder. Su GEOM-DRUG, la costante di Lipschitz empirica è ridotta di 4,6 volte, confermando che la varietà di rappresentazione è più adatta alla modellazione di diffusione.
• Ricchezza della Rappresentazione: Il rango effettivo della rappresentazione aumenta da 65,32 (encoder grezzo) a 88,42 (LENSEs), indicando una riduzione del collasso dimensionale e uno spazio latente più ricco e strutturato.
• Utilità a valle: Quando utilizzato per compiti di regressione a valle sulle proprietà QM9 con encoder congelato e testa di rappresentazione, LENSEs supera l'encoder grezzo su tutte le proprietà, dimostrando che il processo di addestramento generativo agisce di per sé come un paradigma di preaddestramento efficace per affinare le rappresentazioni molecolari.
• Efficienza del Campionamento: LENSEs dimostra un'efficienza di campionamento migliorata, raggiungendo alta validità e stabilità con meno passaggi di generazione rispetto alla baseline.

Significato e Affermazioni
Il documento afferma che la generazione di alta qualità di molecole 3D dipende non solo dall'architettura del generatore ma criticamente dalla qualità della rappresentazione di condizionamento. LENSEs dimostra che lo spazio latente degli encoder preaddestrati può essere sistematicamente migliorato per essere più liscio, più informativo semanticamente e meglio allineato con gli obiettivi generativi.

Gli autori ipotizzano che il loro lavoro stabilisca una relazione bidirezionale tra qualità della generazione e qualità della rappresentazione. Nello specifico, suggeriscono che l'addestramento generativo con obiettivi di allineamento può servire come un potenziale paradigma di preaddestramento per gli encoder molecolari, affinando le rappresentazioni per renderle più predittive e utili per compiti a valle. Il lavoro evidenzia che le rappresentazioni molecolari utili non sono intrinsecamente favorevoli alla generazione e richiedono un affinamento esplicito per sbloccare il loro pieno potenziale nella modellazione generativa.