Riepilogo Tecnico: SinAE – Un Autoencoder di Flow-Matching a Architettura Singola per Sistemi Atomici Cross-Domain

1. Definizione del Problema

La generazione di strutture atomiche 3D è fondamentale per la scoperta di farmaci, la progettazione di materiali e l'ingegneria proteica. Tuttavia, le attuali pipeline generative sono frammentate in tre domini distinti: piccole molecole, cristalli e proteine. Ciascun dominio si affida ad architetture specializzate, adattate ai propri priori strutturali specifici:

• Molecole: Utilizzano tipicamente reti di message-passing equivarianti per gestire la connettività dei legami e le regole di valenza.
• Cristalli: Richiedono la gestione delle condizioni al contorno periodiche e dell'invarianza della base del reticolo.
• Proteine: Spesso si affidano a parametrizzazioni a livello di residuo (es. frame SE(3)) o meccanismi di attenzione a triangolo.

Questa frammentazione esacerba la scarsità di dati, poiché i dataset per singolo dominio (es. QM9 con 134k molecole, MP-20 con 45k cristalli) sono piccoli secondo gli standard moderni. Sebbene l'addestramento cross-domain possa mitigare questo problema, gli approcci esistenti affrontano due ostacoli primari:

1. Limitazioni della Generazione Diretta: Generare direttamente nello spazio delle coordinate 3D costringe un singolo modello a apprendere simultaneamente composizioni valide, geometrie precise e eterogenei priori strutturali (periodicità, backbone di residui, valenza), richiedendo spesso una tokenizzazione complessa e specifica per dominio o una generazione a stadi.
2. Limitazioni degli Autoencoder Latenti: Gli attuali autoencoder latenti cross-domain soffrono di problemi di fedeltà di ricostruzione (i decoder lossy introducono errori geometrici che i priori non possono correggere) e complessità di addestramento (richiedono combinazioni di loss e pesatura specifiche per dominio). Nessun metodo precedente ha raggiunto un autoencoder con obiettivo unificato, architettonicamente generale e privo di perdite simultaneamente per tutti e tre i domini.

2. Metodologia: SinAE

Gli autori introducono SinAE, un autoencoder di flow-matching a architettura singola progettato per molecole, cristalli e proteine.

Architettura

• Backbone Unificato: SinAE utilizza un Transformer vanilla sia per l'encoder che per il decoder. Evita esplicitamente operatori equivarianti, message-passing su grafi o moduli specifici per dominio (es. layer aware del reticolo o attenzione a triangolo).
• Featurizzazione dell'Input:
  • Coordinate: Vengono utilizzate le coordinate cartesiane grezze senza descrittori invarianti (es. distanze pairwise) o costruzione di frame canonici.
  • Periodicità: Per i cristalli, i tre vettori del reticolo sono aggiunti come "token di atomi virtuali" alla sequenza, permettendo al modello di elaborare la matrice del reticolo in modo identico alle coordinate atomiche.
  • Equivarianza: Invece di un'equivarianza architettonica, il modello si affida all'augmentation di rotazione SO(3) e traslazione casuale durante l'addestramento per garantire la robustezza all'orientamento.
• Tokenizzazione: Gli atomi sono rappresentati dalla somma di embedding di elementi appresi e proiezioni lineari delle loro coordinate. Per le proteine, i residui sono trattati come singoli token nelle loro posizioni $C_\alpha$.

Obiettivi di Addestramento

Il modello viene addestrato in due fasi:

1. Flow-Matching Autoencoder:
   • Encoder: Un Transformer stocastico mappa la struttura di input $G=(X, A)$ in un codice latente per ogni token $z$.
   • Decoder: Un Transformer non equivariante ricostruisce la struttura da $z$ utilizzando il flow-matching iterativo. A differenza dei decoder VAE a passaggio singolo, il decoder affina la struttura attraverso molteplici passi ODE.
   • Funzione di Loss: Un obiettivo unificato combina il flow matching condizionale continuo per le coordinate ($L_{coord}$) e il flow matching discreto per il tipo di atomo ($L_{atom}$), oltre a un leggero regolarizzatore KL ($L_{KL}$). Non sono richiesti pesature della loss o loss ausiliarie specifiche per dominio.
2. Latent Flow Matching Prior:
   • Un Diffusion Transformer (DiT) vanilla apprende la distribuzione $p(z)$ sugli output dell'encoder.
   • Poiché il termine KL spinge il posterior verso una Gaussiana standard, il prior opera su uno spazio latente liscio e approssimativamente Gaussiano, senza vincoli di equivarianza.
   • Durante la generazione, un label di dominio viene utilizzato come condizione (tramite classifier-free guidance) per guidare il DiT verso molecole, cristalli o proteine.

3. Contributi Chiave

1. Tokenizzazione a Architettura Singola: SinAE è la prima architettura a gestire molecole, cristalli e backbone proteici utilizzando un singolo Transformer vanilla senza operatori equivarianti o specifici per dominio.
2. Ricostruzione Quasi-Lossless: Spostando il carico della ricostruzione su un decoder di flow-matching iterativo, SinAE raggiunge una RMSD sub-milli-Angström attraverso i domini:
   • Molecole: 0.0002 Å su QM9, 0.0006 Å su GEOM-Drugs.
   • Cristalli: 0.0017 Å su MP-20.
   • Proteine: ~0.013 Å di CA-RMSD del backbone su CASP15.
   • Ciò rappresenta un miglioramento di 10–100× nella fedeltà di ricostruzione rispetto ai precedenti baseline latenti, con tassi di corrispondenza strutturale del 100% su molecole e cristalli.
3. Trasferimento Cross-Domain Positivo: L'addestramento congiunto su QM9 e MP-20 migliora rigorosamente le prestazioni di ricostruzione e generazione in entrambi i domini rispetto all'addestramento a dominio singolo, fornendo prova che un latente atomico condiviso cattura regolarità geometriche trasferibili (es. distribuzioni delle lunghezze di legame, pattern di coordinazione).

4. Risultati Sperimentali

Gli autori hanno valutato SinAE su benchmark standard (QM9, GEOM-Drugs, MP-20, AFDB-FS/CASP15, QMOF).

• Fedeltà di Ricostruzione: SinAE supera tutti i baseline (inclusi GeoLDM, CDVAE, ADiT, UAE-3D, ProteinAE) in termini di RMSD e tassi di corrispondenza strutturale. Il decoder iterativo risolve il posizionamento atomico precocemente e affina la geometria dei legami/sterica successivamente, raggiungendo un'alta precisione senza rilassamento post-hoc.
• Generazione Molecolare: Su QM9 e GEOM-Drugs, il prior DiT latente raggiunge un'alta validità (score PB Valid di 0.97 e 0.93, rispettivamente), superando o eguagliando generatori equivarianti specializzati (es. EQGAT-diff, SemlaFlow).
• Generazione di Cristalli: Su MP-20, il modello SinAE congiunto raggiunge i punteggi S.U.N. (Stable, Unique, Novel) e M.S.U.N. più elevati tra i metodi confrontati, inclusi quelli con architetture invarianti del reticolo. Dimostra una stabilità termodinamica superiore (strutture metastabili) nonostante l'uso di un Transformer generico.
• Generazione Proteica: SinAE raggiunge simultaneamente un'alta progettabilità (0.96 scRMSD $\le$ 2.0 Å) e diversità, superando ProteinAE e LatentDiff. Nei compiti di co-design sequenza-struttura, ottiene punteggi scTM e scRMSD competitivi per lunghezze di 100 e 200.

5. Significato e Rivendicazioni

Il paper sostiene che SinAE dimostri che la generazione atomica cross-domain può essere unificata sotto un'unica, semplice architettura senza sacrificare la fedeltà.

• Semplicità vs Performance: Gli autori sostengono che operatori complessi specifici per dominio non siano strettamente necessari per una generazione di alta qualità se il decoder è sufficientemente potente (flow-matching iterativo) e lo spazio latente è opportunamente regolarizzato.
• Efficienza dei Dati: I risultati forniscono una prova diretta che i latenti atomici condivisi permettono a diverse popolazioni chimiche di beneficiare l'una dell'altra, migliorando le prestazioni in regimi con pochi dati.
• Decoupling di Geometria e Priori: Il lavoro suggerisce che la complessità geometrica può essere gestita dal raffinamento iterativo del decoder, permettendo all'encoder e al prior di operare su uno spazio latente semplificato e domain-agnostic.

Limitazioni Riconosciute:

• La valutazione delle proteine è attualmente limitata alla generazione del backbone; la modellazione all-atom non è ancora stata affrontata.
• L'attuale modello non supporta la generazione condizionata per proprietà, composizione o lunghezza in un singolo checkpoint attraverso tutti i tre domini simultaneamente.
• Catene proteiche più lunghe (>256 residui) richiedono il riaddestramento con budget di lunghezza superiori.