Protein Language Model Embeddings Improve Generalization of Implicit Transfer Operators

Sintesi Tecnica: Gli Embedding dei Modelli di Linguaggio Proteico Migliorano la Generalizzazione degli Operatori di Trasferimento Impliciti

Definizione del Problema

Le simulazioni di dinamica molecolare (MD) sono essenziali per predire osservabili sperimentali da distribuzioni molecolari ad alta dimensionalità, come le distribuzioni di equilibrio di Boltzmann e le densità di transizione. Tuttavia, la MD convenzionale è computazionalmente proibitiva a causa del vasto divario temporale tra i passi di integrazione nell'ordine dei femtosecondi e i tempi di rilassamento biologico nell'ordine dei microsecondi o secondi. Sebbene i metodi di Dinamica Molecolare Generativa (GenMD), come i Boltzmann Generators/Emulators e gli Operatori di Trasferimento Impliciti (ITO), offrano alternative di campionamento efficienti, essi spesso soffrono di una scarsa efficienza dei dati e di una limitata trasferibilità verso sistemi molecolari non visti. Gli approcci esistenti richiedono tipicamente ampie collezioni di lunghe traiettorie MD e faticano a generalizzare a sistemi proteici fuori distribuzione (OOD) senza un fine-tuning specifico per il sistema.

Metodologia

Gli autori introducono PLaTITO (Protein-Language-aware Transferable Implicit Transfer Operator), un framework progettato per migliorare l'efficienza dei dati e la generalizzazione dei surrogati MD incorporando informazioni ausiliarie provenienti da modelli pre-addestrati.

Framework Core: Coarse-Grained TITO

La base del lavoro è un operatore di trasferimento implicito (TITO) a grana grossa (coarse-grained).

• Rappresentazione: Il modello opera su una rappresentazione a grana grossa utilizzando solo le coordinate dello scheletro $C_\alpha$ ($x_t \in \mathbb{R}^{3L}$), dove $L$ è il numero di residui.
• Obiettivo di Apprendimento: Il modello approssima la densità di transizione a lungo termine $p(x_{t+\Delta t} | x_t, \Delta t, S, T)$, dove $S$ è la sequenza amminoacidica e $T$ è la temperatura.
• Meccanismo di Addestramento: Gli autori impiegano il Conditional Flow Matching (CFM). Apprendono un campo di velocità $v_\theta$ che trasforma una variabile latente $z_s$ (interpolando tra il rumore gaussiano e la conformazione futura $x_{t+\Delta t}$) in la distribuzione target. L'obiettivo di addestramento minimizza l'errore quadratico tra la velocità predetta e il campo di velocità condizionale derivato dall'interpolante lineare.
• Architettura: Un'architettura transformer a due stadi separa la rete di condizionamento ($f_c$) dalla rete di predizione della velocità ($f_v$). Entrambe sono transformer non equivarianti basati sull'architettura Proteina, che utilizzano rappresentazioni di residuo e di coppia.

Incorporazione di Rappresentazioni Ausiliarie

Per potenziare la generalizzazione e l'efficienza dei dati, gli autori aumentano gli input di condizionamento con tre fonti di informazioni pre-addestrate:

1. Embedding di Sequenza (pLM): Embedding a livello di residuo estratti da modelli di linguaggio proteico pre-addestrati (specificamente ESM-Cambrian). Essi catturano implicitamente informazioni evolutive, strutturali e funzionali.
2. Embedding di Struttura: Rappresentazioni estratte da un modello consapevole della struttura (Proteina) basato sulle coordinate dello scheletro in input. Queste forniscono caratteristiche geometriche locali e globali.
3. Annotazioni derivate da LLM: Valutazioni euristiche generate da Large Language Models (DeepSeek Reasoner) riguardanti la idoneità di simulare un determinato dominio proteico in isolamento (ad esempio, verificando la presenza di cofattori mancanti o partner di legame).

Il modello finale, PLaTITO-19M, scala il numero di parametri a 19 milioni e utilizza gli embedding del modello più grande ESM-Cambrian 6B.

Contributi Chiave

1. Sviluppo di un TITO Trasferibile: Gli autori hanno sviluppato modelli TITO a grana grossa che generalizzano a sistemi proteici OOD quando addestrati da zero su diverse traiettorie MD fuori equilibrio attraverso molteplici temperature, eliminando la necessità di fine-tuning specifico per il sistema.
2. Integrazione di Segnali Ausiliari: Hanno dimostrato che il condizionamento su embedding di sequenza e struttura pre-addestrati migliora significativamente l'efficienza dei dati. La variante PLaTITO supera costantemente i baseline pur richiedendo sostanzialmente meno dati e risorse computazionali.
3. Prestazioni allo Stato dell'Arte: PLaTITO-19M raggiunge risultati allo stato dell'arte nei benchmark di campionamento di equilibrio per sistemi proteici OOD, inclusi i processi di ripiegamento rapido (fast-folding), superando il baseline BioEmu.
4. Cinetica Non-Arrhenius: Lo studio mostra che la cinetica dipendente dalla temperatura appresa da PLaTITO esibisce un comportamento non-Arrhenius, coerente con i paesaggi di energia libera complessi e accidentati del ripiegamento proteico.

Risultati

I modelli sono stati valutati sul dataset mdCATH (traiettorie diverse fuori equilibrio) e testati su proteine a ripiegamento rapido (es. Villin, dominio WW, A3D) e tasche di legame criptiche.

• Campionamento di Equilibrio: PLaTITO-19M ha ottenuto un Mean Absolute Error (MAE) di 0.824 e un Root Mean Squared Error (RMSE) di 1.099 sui paesaggi di energia libera, superando BioEmu (MAE 1.110, RMSE 1.389) e altri baseline come ConfDiff e Str2Str.
• Efficienza di Dati e Calcolo: PLaTITO-19M ha raggiunto questi risultati con 1.100 ore di GPU e 56 ms di dati MD aggregati. Al contrario, BioEmu ha richiesto 9.216 ore di GPU e 216 ms di dati, oltre a un ulteriore addestramento su 131k strutture AFDB. Anche varianti TITO più piccole (3M parametri) hanno eguagliato le prestazioni di BioEmu con molte meno risorse.
• Dinamica a Lungo Termine: Il modello ha generato con successo traiettorie da 120 $\mu$s per A3D, catturando eventi ripetuti di ripiegamento e srotolamento (unfolding). I tempi medi di primo passaggio (MFPT) stimati erano coerenti con i trend della MD, sebbene sistematicamente più veloci (un noto effetto variazionale).
• Dipendenza dalla Temperatura: Il modello ha recuperato tassi di ripiegamento e srotolamento non-Arrhenius attraverso le temperature (320K–440K), indicando che ha appreso una dinamica fisicamente significativa oltre la semplice scalatura esponenziale.
• Tasche Criptiche: PLaTITO-19M ha dimostrato una migliore copertura nel campionamento di tasche di legame criptiche rispetto a BioEmu, recuperando con successo stati "apo" in casi difficili in cui BioEmu falliva.

Significato e Rivendicazioni

Il paper sostiene che sfruttare la struttura di auto-correlazione dei dati MD attraverso l'apprendimento ITO, combinato con il condizionamento ausiliario da modelli di linguaggio proteico pre-addestrati, offre un percorso altamente efficiente in termini di dati per apprendere i surrogati della MD.

• Generalizzazione: La scoperta principale è che gli embedding pLM forniscono potenti bias induttivi che permettono ai modelli TITO di apprendere la trasferibilità in modo più efficiente da traiettorie diverse fuori equilibrio, consentendo prestazioni su sistemi OOD senza fine-tuning.
• Efficienza: Il lavoro evidenzia una disparità centrale: mentre i Boltzmann Emulators si affidano a enormi quantità di dati per mappare le densità di equilibrio, i modelli TITO sfruttano le correlazioni temporali per abbassare i costi di dati e calcolo di quasi un ordine di grandezza mantenendo o superando le prestazioni.
• Fedeltà Fisica: La capacità di catturare la cinetica non-Arrhenius ed esplorare tasche criptiche suggerisce che il modello apprenda il propagatore sottostante del sistema piuttosto che una semplice distribuzione statica.

Gli autori riconoscono i limiti, tra cui l'uso di rappresentazioni a grana grossa (ignorando i riarrangiamenti delle catene laterali) e la mancanza di garanzie formali riguardo al dettaglio bilanciato o alla stabilità. Notano che, sebbene le prestazioni siano elevate, è necessario aumentare ulteriormente la scala dei dati e la capacità del modello per catturare tutte le fasi metastabili e migliorare l'accuratezza delle scale temporali.