UBio-MolFM: A Universal Molecular Foundation Model for Bio-Systems

UBio-MolFM è un modello fondazionale universale per i sistemi biologici che, grazie a un nuovo dataset bio-specifico, un'architettura transformer equivariante scalabile e un protocollo di apprendimento curricolare, supera il compromesso tra precisione quantistica e scala biologica, consentendo simulazioni molecolari all-atom di sistemi complessi con fedeltà pari a quella dei calcoli *ab initio*.

L'essenziale

Immagina di voler studiare come funziona una macchina complessa, come un'auto da corsa o un orologio svizzero. Hai due modi per farlo:

1. Il Microscopio Quantistico (QM): È come smontare ogni singola vite, ingranaggio e molla per vedere esattamente come sono fatti e come si muovono a livello atomico. È incredibilmente preciso, ma ci vorrebbe un'eternità per farlo su un'auto intera. È troppo lento e costoso per sistemi grandi.
2. La Macchina Giocattolo (MM): È come guardare l'auto da lontano. È veloce, puoi guidarla e vedere come si comporta, ma non sai esattamente cosa succede dentro ogni ingranaggio. È veloce, ma imprecisa.

Per decenni, gli scienziati hanno dovuto scegliere tra precisione e velocità. Non potevano avere entrambe.

UBio-MolFM è la soluzione magica che il team di ricerca IQuest ha creato per rompere questo compromesso. È come se avessero costruito un "Super Occhio Digitale" che vede ogni singolo atomo con precisione quantistica, ma che lavora alla velocità di un videogioco.

Ecco come funziona, spiegato con tre ingredienti principali:

1. L'Alimentazione: Il "Libro di Ricette" Gigante (UBio-Mol26)

Per insegnare a un'intelligenza artificiale a fare previsioni, le servono milioni di esempi. Prima, gli scienziati avevano libri di ricette con solo piccoli ingredienti (molecole piccole).

• La novità: Hanno creato un nuovo libro di ricette, UBio-Mol26, che contiene "piatti" molto più grandi e complessi: proteine, DNA e RNA, proprio come quelli che esistono nel nostro corpo.
• La strategia "Due Forchette": Hanno usato due metodi per raccogliere questi dati:
  • Dal basso verso l'alto: Hanno costruito milioni di piccoli pezzi (come mattoncini Lego) per capire le regole di base.
  • Dall'alto verso il basso: Hanno preso foto di proteine reali (come se fossero edifici) e ne hanno tagliato dei pezzi per studiare come si comportano nel loro ambiente naturale (con l'acqua intorno).

2. Il Cervello: Il "Motore Leggero" (E2Former-V2)

Anche con un libro di ricette perfetto, il cervello dell'AI deve essere veloce. I modelli precedenti erano come camion pesanti: potenti, ma lenti e ingombranti.

• La novità: Hanno creato E2Former-V2, un motore che è come una F1 elettrica.
• Il trucco: Invece di guardare tutti gli atomi e calcolare le loro interazioni con tutti gli altri (cosa che farebbe esplodere la memoria del computer), questo modello usa un sistema intelligente di "filtri". Guarda da vicino i vicini immediati e da lontano solo le cose importanti, ignorando il rumore di fondo.
• Risultato: È fino a 4 volte più veloce dei migliori modelli attuali, permettendo di simulare sistemi enormi (fino a 1.500 atomi) senza bloccarsi.

3. La Scuola: L'allenamento a "Livelli" (Curriculum Learning)

Non puoi insegnare a un bambino a fare l'equazione di una fisica quantistica prima di avergli insegnato a contare.

• La novità: Hanno usato un metodo di insegnamento a tre livelli:
  1. Livello 1: Impara le basi con molecole piccole e semplici (veloce).
  2. Livello 2: Impara a collegare l'energia alle forze (come spingere un oggetto), assicurandosi che la fisica sia coerente.
  3. Livello 3: Si tuffa nei sistemi biologici complessi (proteine, acqua, ioni) per affinare la precisione.
• Questo ha permesso al modello di non "confondersi" quando passa da molecole piccole a quelle giganti.

Cosa ha scoperto? (I Risultati)

Hanno messo alla prova il loro "Super Occhio" su scenari reali:

• L'Acqua: Ha simulato l'acqua liquida e il sale disciolto con una precisione che corrisponde quasi perfettamente alla realtà sperimentale.
• Le Proteine: Ha visto come una proteina cambia forma a seconda che sia nell'acqua o nel vuoto (come un origami che si apre e si chiude).
• Il DNA e i Metalli: Ha simulato come gli ioni di magnesio si attaccano al DNA, un dettaglio cruciale per capire come funzionano i geni, con una precisione che i modelli precedenti non avevano.

In sintesi

UBio-MolFM è come aver dato agli scienziati un microscopio che non si blocca mai. Prima, dovevano scegliere tra vedere i dettagli (ma essere lentissimi) o vedere il quadro generale (ma essere imprecisi). Ora, con questo strumento, possono vedere i dettagli atomici di sistemi biologici enormi in tempi ragionevoli.

È un passo enorme verso la "biologia eseguibile", dove potremo simulare malattie, progettare farmaci e capire la vita stessa direttamente al computer, con una precisione che prima era impossibile.

Sommario tecnico

1. Il Problema: Il Divario Scala-Accuratezza

La simulazione di sistemi biologici con precisione quantomeccanica (QM) rappresenta una sfida fondamentale nelle scienze computazionali. Esiste un persistente "divario scala-accuratezza":

• Metodi Ab Initio (QM): Forniscono la precisione elettronica necessaria per fenomeni come polarizzazione, trasferimento di carica e chimica reattiva, ma la loro complessità computazionale scala cubicamente ($O(N^3)$ o superiore), limitandoli a sistemi di poche centinaia di atomi.
• Meccanica Molecolare Classica (MM): Scalabile a milioni di atomi, ma utilizza forme funzionali fisse che faticano a rappresentare le complesse superfici di energia potenziale (PES) delle macchine biologiche.
• Campi di Forza basati su ML (MLFF) esistenti: Le soluzioni attuali sono limitate per i sistemi mesoscopici biologici a causa di:
  1. Copertura dei dati insufficiente: I dataset pubblici (es. SPICE, OMol25) si concentrano su piccole molecole o sono limitati a sistemi di circa 350 atomi.
  2. Architetture limitate: Molti modelli si basano su cutoff locali, sottorappresentando le elettrostatiche a lungo raggio e causando errori di consistenza dimensionale.
  3. Costo computazionale: I modelli equivarianti di alto ordine (es. MACE, NequIP) sono troppo pesanti per traiettorie MD lunghe di proteine solvate.

2. Metodologia: Il Framework UBio-MolFM

UBio-MolFM è un framework di fondazione universale progettato per colmare questo divario attraverso tre innovazioni sinergiche:

A. Costruzione dei Dati: UBio-Mol26

È stato creato un dataset specifico per la biologia contenente 17 milioni di configurazioni con sistemi fino a 1.200 atomi. La strategia di costruzione è "a due punte":

• Bottom-Up: Enumerazione combinatoria sistematica di blocchi fondamentali (es. tutti i 6.840 possibili tripeptidi, basi DNA/RNA) per garantire una copertura imparziale dello spazio chimico.
• Top-Down: Campionamento diretto da ambienti proteici nativi (utilizzando AlphaFold Protein Structure Database), estrazione di cluster sferici centrati su residui specifici con acqua esplicita, per massimizzare la rilevanza biologica.
• Metodologia Computazionale: Utilizzo del funzionale $\omega$B97M-D3 con una strategia di basi miste (def2-TZVP per H e metalli, def2-TZVPD per altri elementi) per bilanciare accuratezza e convergenza SCF su sistemi grandi. Una parte dei dati utilizza def2-SVP per aumentare la scala dei dati a costi computazionali ridotti.

B. Architettura del Modello: E2Former-V2

Un transformer equivariante a scala lineare che integra due componenti chiave:

1. EAAS (Equivariant Axis-Aligned Sparsification): Una riduzione algebrica che converte i prodotti tensoriali densi SO(3) in operazioni sparse tramite re-indicizzazione SO(2), mantenendo l'equivarianza esatta ma riducendo drasticamente il costo computazionale (circa 6x più veloce).
2. Modellazione Long-Short Range (LSR): Un'architettura ibrida che combina:
   • Un modulo a corto raggio (raggio ~5 Å) per le interazioni locali ad alta risoluzione.
   • Un modulo a lungo raggio (raggio ~15 Å) basato su un grafo bipartito atomo-frammento per catturare elettrostatiche e correlazioni a distanza senza creare un grafo completamente connesso.
3. Efficienza Hardware: Implementazione di kernel GPU personalizzati (Triton) che calcolano l'attenzione "on-the-fly" senza materializzare esplicitamente le matrici di bordo, riducendo l'uso di memoria HBM e permettendo la scalabilità a sistemi di 100.000 atomi.

C. Protocollo di Addestramento: Curriculum Learning a Tre Stadi

Per gestire l'eterogeneità dei dati (dalle piccole molecole ai grandi complessi biologici), è stato adottato un approccio a stadi:

1. Stadio 1 (Inizializzazione Rapida): Addestramento su OMol25 (piccole molecole) con testate separate per energia e forza per una rapida convergenza.
2. Stadio 2 (Coerenza Energia-Forza): Rimozione della testata forza indipendente; le forze sono calcolate come gradienti dell'energia ($F = -\nabla E$) per garantire la conservazione dell'energia e la consistenza fisica della PES.
3. Stadio 3 (Affinamento Multi-Fidelità): Integrazione dei dati UBio-Mol26 (SVP e TZVPD) con un tasso di apprendimento ridotto. Utilizzo di una strategia di supervisione focalizzata sulla forza per mitigare gli offset energetici tra diversi livelli di teoria (funzionali e basi diverse).

3. Risultati Chiave

Accuratezza Microscopica e Generalizzazione Out-of-Distribution (OOD)

• Il modello è stato testato su sistemi 1.300–1.500 atomi (più grandi del training set massimo di 1.200).
• UBio-MolFM (S3) ha superato significativamente i modelli di base (MACE-OMol, UMA-S-1p1) negli errori di forza e nelle differenze di energia relativa.
• In particolare, per le proteine e l'RNA, l'errore MAE dell'energia è diminuito drasticamente (es. >60% di miglioramento sull'RNA rispetto ai baseline).
• Il modello dimostra una capacità superiore nel tracciare l'evoluzione temporale della PES, cruciale per MD stabili.

Fidelity Fisica nelle Simulazioni MD

• Struttura di Solvatazione: Riproduce con alta fedeltà la struttura dell'acqua liquida e la solvatazione ionica (NaCl), corrispondendo ai dati sperimentali e DFT per le funzioni di distribuzione radiale (RDF) e i numeri di coordinazione.
• Dinamica Proteica (Ciclossporina A): Il modello mantiene correttamente le conformazioni dipendenti dal solvente (stato "aperto" in acqua vs "chiuso" nel vuoto), dimostrando una corretta competizione tra legami idrogeno intramolecolari e interazioni solvente-soluto.
• Dinamica RNA e Ioni Metallici: Nel sistema RNA 1L2X con $Mg^{2+}$, UBio-MolFM è l'unico modello che riproduce accuratamente sia le distanze Mg-O (~2.04 Å) che gli angoli flessibili Mg-O-P, superando le limitazioni dei campi di forza classici (Amber99) e di altri modelli ML (UMA).

Efficienza di Inferenza

• Su sistemi di grandi dimensioni (1K–50K atomi), UBio-MolFM raggiunge un throughput di ~4 volte superiore rispetto ai modelli equivarianti più avanzati (es. UMA-S) su una singola GPU H100.
• Grazie all'architettura EAAS e LSR, il modello scala fino a 50.000 atomi, mentre molti baseline falliscono per mancanza di memoria (OOM) già a 10.000 atomi.

4. Contributi e Significato

• Ponte tra QM e Scala Biologica: UBio-MolFM dimostra che è possibile ottenere fedeltà a livello ab initio su sistemi biologici complessi (fino a ~1.500 atomi) mantenendo un throughput pratico per le simulazioni MD.
• Dataset Biologico Su Misura: UBio-Mol26 colma il vuoto nei dati pubblici, fornendo la prima risorsa su larga scala specifica per macromolecole biologiche solvate, superando i limiti di OMol25.
• Efficienza Hardware: L'introduzione di EAAS e kernel di attenzione "edge-free" risolve i colli di bottiglia di memoria tipici dei transformer equivarianti, rendendo fattibile l'uso di questi modelli su sistemi di scala reale.
• Open Science: Il team ha pianificato il rilascio pubblico dei pesi del modello, del codice di inferenza e di un subset del dataset (UBio-Protein26 5M) per promuovere la ricerca collaborativa.

Conclusione

UBio-MolFM non è presentato come una soluzione definitiva, ma come una prova di concetto robusta che la simulazione biologica ad alta fedeltà e generalizzabile è raggiungibile. Il lavoro segna un passo significativo verso una "biologia eseguibile", dove simulazioni quantomeccaniche accurate diventano strumenti di routine per decifrare i meccanismi molecolari della vita. Le sfide future includono il miglioramento della stabilità temporale per il DNA e l'ulteriore ottimizzazione hardware per ridurre il divario di velocità rispetto ai campi di forza classici.