Robust Inference-Time Steering of Protein Diffusion Models via Embedding Optimization

Riepilogo Tecnico: Steering Robusto a Tempo di Inferenza di Modelli di Diffusione Proteica tramite Ottimizzazione degli Embedding

Enunciato del Problema
Una sfida centrale nella biofisica strutturale è generare conformazioni biomolecolari che soddisfino sia il realismo fisico che le misurazioni sperimentali. Sebbene i modelli di diffusione da sequenza a struttura (ad esempio AlphaFold 3, Protenix) forniscano potenti prior guidati dai dati, la risoluzione di problemi inversi di determinazione strutturale richiede spesso il recupero di stati situati in regioni a bassa densità della distribuzione a priori. Gli approcci attuali, come il Campionamento Posteriore di Diffusione (DPS), guidano la generazione perturbando le coordinate atomiche utilizzando gradienti derivati dalle verosimiglianze sperimentali. Tuttavia, quando l'obiettivo si trova lontano dalle modalità ad alta probabilità del prior, il DPS richiede un'aggressiva sovrappesatura della verosimiglianza. Ciò porta a un paesaggio di campionamento mal condizionato, causando instabilità, sensibilità agli iperparametri e incapacità di sfuggire agli ottimi locali, in particolare in compiti come il fitting di mappe di microscopia elettronica criogenica (cryo-EM).

Metodologia: EmbedOpt
Gli autori propongono EmbedOpt, un framework di steering a tempo di inferenza che introduce un asse di ottimizzazione ortogonale ai metodi basati sulle coordinate. Invece di perturbare le coordinate atomiche rumorose $x_t$ sotto un prior fisso, EmbedOpt ottimizza direttamente l'embedding condizionale $c$ del modello.

• Meccanismo: Nelle architetture in stile AlphaFold 3, l'embedding condizionale codifica ricchi segnali co-evolutivi derivati da allineamenti di sequenze multiple (MSA). EmbedOpt tratta il problema inverso come un compito di massimizzazione della ricompensa, dove la ricompensa $R(x_0)$ corrisponde alla verosimiglianza sperimentale.
• Algoritmo: Ad ogni passo di diffusione $t$, EmbedOpt esegue un singolo passo di ascesa del gradiente sull'embedding condizionale $c_t$ per massimizzare una ricompensa surrogata valutata sulla predizione denoizzata $\hat{x}_\theta(x_t, c_t, \sigma_t)$. L'embedding aggiornato $c_{t-1}$ viene quindi utilizzato per avanzare la traiettoria di diffusione.
• Base Teorica: Gli autori forniscono un'analisi teorica che stabilisce che, sotto aggiornamenti a piccoli passi e condizioni di regolarità adeguate, EmbedOpt agisce come un metodo di regione di fiducia locale. Ciò garantisce un miglioramento monotono della ricompensa surrogata, evitando le sensibilità geometriche che affliggono gli aggiornamenti basati sulle coordinate.
• Confronto con DPS: Mentre DPS applica i gradienti di verosimiglianza direttamente allo spazio delle coordinate ($x_t$), accoppiando l'aggiornamento alla sensibilità del denoiser al rumore, EmbedOpt applica un precondizionatore nello spazio degli embedding. Ciò permette alla direzione di aggiornamento delle coordinate di rimanere all'interno dello span del gradiente di ricompensa, pur essendo modulata dalla geometria dello spazio degli embedding, risultando in dinamiche di ottimizzazione più fluide.

Contributi Chiave

1. Nuovo Asse Algoritmico: EmbedOpt propone l'ottimizzazione dell'embedding della sequenza condizionale come meccanismo primario per guidare i modelli di diffusione proteica, complementando l'approccio dominante nello spazio delle coordinate.
2. Analisi Teorica: Il lavoro stabilisce una caratterizzazione di regione di fiducia che prova il comportamento di ottimizzazione monotona per EmbedOpt sotto aggiornamenti a piccoli passi.
3. Robustezza ed Efficienza: EmbedOpt dimostra prestazioni stabili su un intervallo di velocità di apprendimento che copre due ordini di grandezza (da 0,01 a 1,0), superando significativamente il DPS nella robustezza agli iperparametri. Raggiunge inoltre risultati comparabili o superiori con meno passi di diffusione (ad esempio, una riduzione di 4 volte dei passi sui vincoli di distanza).
4. Validazione nel Mondo Reale: Il metodo è stato validato su dati sperimentali reali e rumorosi di cryo-EM, mostrando prestazioni competitive o superiori rispetto alla baseline pubblicata CryoBoltz, mantenendo al contempo una geometria stereo-chimica significativamente migliore.

Risultati Sperimentali
Gli autori hanno valutato EmbedOpt contro DPS e la baseline CryoBoltz su tre benchmark:

1. Vincoli di Distanza Sparsi: Su 24 proteine multi-dominio, EmbedOpt ha eguagliato il DPS nelle prestazioni di picco, ma ha mantenuto geometrie valide e alti tassi di soddisfazione dei vincoli con 4 volte meno passi di diffusione. Ha mostrato una traccia di ricompensa fluida e monotona, mentre DPS ha mostrato alta volatilità ad alta frequenza.
2. Fitting di Mappe Cryo-EM Sintetiche: Su 77 sistemi, EmbedOpt ha superato il DPS su target "difficili" dove il prior e la mappa sperimentale erano disallineati. EmbedOpt ha risolto con successo riarrangiamenti globali dei domini che hanno causato il ristagno del DPS in ottimi locali.
3. Fitting Cryo-EM Sperimentale Reale: Su 6 target dal benchmark CryoBoltz, EmbedOpt (utilizzando il prior Protenix) ha superato la baseline CryoBoltz pubblicata (utilizzando Boltz-1) su 5 sistemi su 6 per quanto riguarda il fitting dei dati (CC mask), la plausibilità fisica (punteggio MolProbity) e l'accordo con il riferimento (RMSD). Notabilmente, mentre CryoBoltz ha raggiunto un'alta correlazione con la mappa, ha spesso prodotto strutture con geometria degradata che il rilassamento energetico non è riuscito a riparare completamente; EmbedOpt ha mantenuto nativamente una stereo-chimica di alta qualità.

Significato e Limitazioni
Il lavoro afferma che EmbedOpt fornisce un'alternativa robusta, efficiente e teoricamente fondata allo steering basato sulle coordinate per i problemi inversi proteici. Spostando la distribuzione a priori tramite ottimizzazione degli embedding invece di forzare il campionamento posteriore a superare un prior non corrispondente, il metodo raggiunge una maggiore stabilità e richiede meno sintonizzazione degli iperparametri.

Gli autori riconoscono modestamente le limitazioni:

• Copertura della Varietà: EmbedOpt è limitato dalla varietà strutturale del modello pre-addestrato. Se la conformazione target è scarsamente rappresentata nei dati di addestramento (ad esempio, a causa di informazioni evolutive limitate o target fuori distribuzione), né EmbedOpt né DPS possono recuperarla.
• Deviazione dalla Varietà: A velocità di apprendimento estremamente aggressive, EmbedOpt può deviare dalla varietà strutturale, sebbene questo confine sia significativamente più lontano rispetto al DPS.
• Ottimi Locali: Sebbene EmbedOpt eviti tipicamente gli ottimi locali nello spazio delle coordinate dove fallisce il DPS, può rimanere intrappolato in ottimi locali nello spazio degli embedding. Gli autori suggeriscono che schemi ibridi che combinano la ricerca globale stabile di EmbedOpt con l'esplorazione locale stocastica del DPS possano essere una direzione futura promettente.

Il codice è disponibile all'indirizzo https://github.com/rs-station/embedopt.