Riepilogo Tecnico: NEAT – Trasformatore Autoregressivo di Insiemi Guidato dal Vicinato, Efficiente per la Generazione Molecolare 3D

Enunciato del Problema

I flussi di lavoro computazionali attuali per la scoperta di molecole bioattive si basano fortemente sullo screening virtuale o sulla generazione de novo. Sebbene lo screening virtuale sia limitato dalla dimensione finita delle librerie chimiche e dai crescenti costi computazionali, la modellazione generativa offre un'alternativa promettente. Tuttavia, la generazione di strutture molecolari 3D presenta sfide uniche. Le molecole sono intrinsecamente tridimensionali e le loro proprietà sono spesso determinate dalla conformazione.

I modelli di generazione 3D all'avanguardia esistenti utilizzano prevalentemente tecniche di diffusione e flow-matching. Questi metodi soffrono spesso di una scalatura quadratica negli aggiornamenti del passaggio di messaggi, mancano di adattabilità intrinseca a dimensioni variabili delle molecole (richiedendo spesso un numero fisso di atomi determinato all'inizializzazione) e faticano a gestire l'inserimento dinamico di atomi durante la denoising.

I modelli autoregressivi offrono un'alternativa flessibile che supporta naturalmente numeri variabili di atomi e richiede una sola chiamata all'encoder per token. Tuttavia, i trasformatori autoregressivi standard richiedono un ordinamento sequenziale dei token. Nei grafi molecolari, gli atomi non hanno un ordine sequenziale intrinseco. Lavori precedenti hanno affrontato questo problema utilizzando ordinamenti canonici degli atomi, ma questi approcci non sono invarianti per permutazione rispetto agli atomi. Introducono un pregiudizio nella previsione del token successivo verso convenzioni di ordinamento specifiche, limitando la capacità del modello di generalizzare o eseguire compiti come il completamento del prefisso, dove l'ordine iniziale degli atomi potrebbe non allinearsi con l'ordine canonico della molecola finale. Al contrario, gli approcci basati su Reti Neurali su Grafi (GNN) che garantiscono l'invarianza per permutazione spesso producono molecole con bassa validità chimica.

Metodologia

Gli autori propongono NEAT (Trasformatore Autoregressivo di Insiemi Guidato dal Vicinato, Efficiente), il primo modello trasformatore autoregressivo per la generazione molecolare 3D che rispetta l'invarianza per permutazione degli atomi senza fare affidamento su ordinamenti canonici degli atomi.

Architettura Principale

NEAT tratta i grafi molecolari come insiemi di atomi piuttosto che come sequenze. L'architettura è composta da tre componenti principali:

1. Modulo di Augmentation dei Dati (Guida dal Vicinato): Invece di codificazioni posizionali sequenziali, il modello utilizza una strategia di campionamento per imitare la crescita molecolare. Dato un grafo molecolare, viene campionato un sottografo connesso (insieme sorgente). L'insieme target è definito come i "nodi di confine" (vicini a 1-hop) dell'insieme sorgente. Questo fornisce un insieme di continuazioni di token valide per il passo successivo, indipendentemente dall'ordine in cui gli atomi sono stati aggiunti.
2. Tronco del Trasformatore di Insiemi: L'insieme sorgente viene codificato utilizzando un trasformatore di insiemi (basato su NanoGPT/GPT-2 con attenzione bidirezionale e senza incorporamenti posizionali). I tipi di atomi sono incorporati linearmente e le posizioni sono codificate utilizzando caratteristiche di Fourier. Un pooling additivo aggrega queste rappresentazioni dei nodi in una rappresentazione a livello di insieme ($z_n$). Questa progettazione garantisce l'invarianza per permutazione rispetto all'ordine degli atomi nell'insieme sorgente.
3. Testa Flow Matching: La rappresentazione a livello di insieme $z_n$ viene utilizzata per prevedere il token successivo, fattorizzata in tipo di atomo e posizione: $p(a_{n+1}, x_{n+1} | z_n) = p(a_{n+1} | z_n) \cdot p(x_{n+1} | a_{n+1}, z_n)$.
   • Tipo di Atomo: Modellato come una distribuzione categorica tramite un livello lineare e softmax.
   • Posizione dell'Atomo: Modellata tramite Flow Matching Condizionale (CFM). Il modello impara un campo vettoriale per trasportare una distribuzione iniziale di posizioni casuali alle posizioni atomiche reali.

Strategia di Addestramento

• Guida dal Vicinato: Durante l'addestramento, il modello non è supervisionato da un singolo token successivo, ma dall'insieme di tutti i nodi di confine validi. Questo elimina la necessità di canonizzazione.
• Teacher Forcing: La previsione del tipo di atomo e il recupero della posizione sono disaccoppiati. Il modello prevede il tipo, e la testa di flow matching impara la posizione condizionata a quel tipo.
• Augmentation: Per migliorare la robustezza contro l'accumulo di errori, viene iniettato rumore gaussiano in una frazione degli atomi dell'insieme sorgente durante l'addestramento, e le posizioni vengono ricalibrate rispetto all'insieme sorgente rumoroso.
• Funzione di Perdita: La perdita totale è la somma della perdita di entropia incrociata per i tipi di atomi e della perdita di flow matching per le posizioni. Viene risolto un problema di assegnazione lineare per accoppiare le distribuzioni sorgente e target, minimizzando i costi di trasporto ed evitando percorsi incrociati.

Inferenza

La generazione inizia con un atomo iniziale casuale. Ad ogni passo, l'insieme corrente di atomi viene codificato, viene previsto il tipo del prossimo atomo e viene campionata una posizione iniziale casuale. Il campo vettoriale di flow matching viene quindi integrato (utilizzando metodi di Eulero o Eulero-Maruyama) per determinare le coordinate 3D finali. Questo processo si ripete fino alla previsione di un token di stop.

Contributi Chiave

1. Generazione Autoregressiva Invariante per Permutazione: NEAT introduce una strategia di addestramento innovativa che utilizza la guida dal vicinato e i trasformatori di insiemi per realizzare una generazione molecolare 3D invariante rispetto all'ordinamento degli atomi, superando i limiti dell'ordinamento canonico nei modelli autoregressivi.
2. Efficienza: Trattando l'input come un insieme e utilizzando una singola chiamata all'encoder per token, NEAT evita la scalatura quadratica del passaggio di messaggi basato sulla diffusione e la necessità di più passi di integrazione per ogni passo di generazione di atomi presenti in alcune basi di riferimento di flow-matching.
3. Capacità di Completamento del Prefisso: La natura agnostica rispetto all'ordine di NEAT consente un completamento del prefisso di alta qualità (generazione di molecole attorno a un dato scaffold) senza richiedere che il prefisso si allinei con un ordinamento canonico specifico, un compito in cui i modelli basati su sequenza come QUETZAL faticano.
4. Prestazioni All'Avanguardia: Il modello ottiene risultati competitivi o superiori su benchmark standard (QM9 e GEOM-Drugs) rispetto sia alle basi di riferimento autoregressive che a quelle di diffusione/flow-matching.

Risultati Sperimentali

Il modello è stato valutato sui dataset QM9 (piccole molecole organiche) e GEOM-Drugs (molecole simili a farmaci).

• Qualità della Generazione: Su QM9, NEAT ha superato i precedenti modelli autoregressivi (G-SchNet, G-SphereNet, Symphony, Geo2Seq) in termini di validità e unicità xyz2mol. Ha raggiunto una validità xyz2mol del 95,2% e un'unicità del 91,6%, classificandosi secondo per validità e unicità di ricerca. Su GEOM, NEAT ha raggiunto la più alta validità di lookup (99,7%) e la seconda migliore validità xyz2mol (93,9%) e unicità (91,6%) tra i modelli autoregressivi, eguagliando le prestazioni di modelli di flow-matching all'avanguardia come FlowMol3.
• Completamento del Prefisso: In un compito di generazione di 1.000 molecole da 100 prefissi diversi, NEAT ha superato significativamente QUETZAL. NEAT ha raggiunto una validità di lookup del 97,8% e una validità xyz2mol del 92,7%, mentre QUETZAL ha raggiunto solo il 27,8% e il 28,1% rispettivamente. Questo dimostra la capacità di NEAT di generalizzare le caratteristiche strutturali indipendentemente dalle convenzioni di ordinamento degli atomi.
• Efficienza Temporale: NEAT dimostra significativi vantaggi di velocità. Esegue circa 5 volte più velocemente di SemlaFlow e 16 volte più velocemente di FlowMol3. Nei casi peggiori, la complessità di NEAT è $O(BN^3)$ (dove $N$ è il numero di atomi), mentre i modelli di diffusione/flow scalano come $O(BTN^2)$ (dove $T$ è il numero di passi di integrazione). Poiché $N$ è tipicamente inferiore a $T$ per le molecole simili a farmaci, NEAT offre un vantaggio pratico in termini di velocità.

Significato e Affermazioni

Il documento afferma che NEAT fornisce una solida base per sfruttare le architetture di trasformatori in compiti complessi di progettazione di farmaci, specificamente per la generazione condizionale di molecole 3D. Rimuovendo la dipendenza dagli ordinamenti canonici degli atomi, il modello abilita compiti flessibili come la generazione basata su scaffold (completamento del prefisso) che erano precedentemente difficili per gli approcci autoregressivi.

Gli autori notano che, sebbene l'architettura attuale non sia esplicitamente equivariante SE(3), l'uso di augmentation con rotazioni casuali durante l'addestramento compensa efficacemente questo aspetto, risultando in una varianza minima nelle previsioni. Il lavoro suggerisce che l'integrazione di NEAT con un trasformatore equivariante SE(3) è una direzione promettente per la ricerca futura.

Gli autori sottolineano che, sebbene il modello generi strutture 3D di alta qualità, la generazione di molecole è solo un passo in una pipeline complessa che richiede estese validazioni, sintesi e test. Riconoscono anche le limitazioni, come la dipendenza dalle informazioni sugli spigoli per la guida dal vicinato (che potrebbe limitare l'applicabilità ai cristalli inorganici) e la mancanza di estrapolazione oltre il numero massimo di atomi osservato durante l'addestramento.