VFUSE: Virulent Feature Understanding with Sparse autoEncoders

Sintesi Tecnica: VFUSE – Comprensione di Caratteristiche Virulente tramite Autoencoder Sparsi

Problematica

I modelli generativi per il design proteico, in particolare i modelli di diffusione all-atom come RoseTTAFold3 (RF3) e RFDiffusion3 (RFD3), possiedono la capacità di doppio uso di progettare proteine pericolose (ad esempio, neurotossine, proteine di evasione immunitaria) scaffoldando nuove strutture attorno a template pericolosi. Gli attuali metodi di screening della biosicurezza si affidano all'omologia di sequenza o a classificatori a livello di sequenza. Questi approcci soffrono di due criticità fondamentali: operano post-hoc (segnalando le sequenze dopo la generazione) e mancano di interpretabilità meccanicistica, il che significa che non possono spiegare quali specifiche caratteristiche strutturali abbiano attivato un segnale di pericolo né intervenire durante il processo di campionamento. Inoltre, sebbene gli Autoencoder Sparsi (SAE) siano stati applicati con successo ai modelli di linguaggio proteico (PLM) e ai modelli di visione, non sono ancora stati utilizzati per sondare le rappresentazioni interne dei modelli di diffusione all-atom.

Metodologia

Gli autori introducono VFUSE, un framework di interpretabilità meccanicistica che addestra Autoencoder Sparsi sulle attivazioni di modelli diffusion-transformer per auditare la presenza di caratteristiche associate a pericoli.

1. Modello e Raccolta delle Attivazioni:

   • Lo studio prende di mira RFDiffusion3 (un modello di diffusione all-atom a 18 blocchi) e RoseTTAFold3 (il suo corrispondente di folding).
   • Le attivazioni ($d=768$) vengono catturate tramite PyTorch hooks in specifici blocchi dello stream residuo: blocchi 6, 8 e 12 per RFD3, e blocchi 12 e 16 per RF3.
   • Per RFD3, gli input sono generati tramite diffusione parziale (a $\partial t = 5$, circa 5Å di rumore) su strutture PDB di template pericolosi per simulare il processo di design del modello. Gli input di RF3 vengono ripiegati direttamente.

2. Costruzione del Dataset:

   • Sequenze Virulente: Tratte da SafeProtein (275 coppie) e ToxinPred3 (1.200 sequenze solo per RF3).
   • Sequenze Benigne: Sequenze con lunghezza corrispondente (100–300 aa) da UniProt, filtrate per escludere parole chiave relative a tossine/virulenza.
   • Controllo dell'Omologia: Per evitare che i probe memorizzino l'identità della famiglia di fold invece di apprendere i segnali di virulenza, le sequenze sono raggruppate per identità al 30% utilizzando mmseqs2. Le divisioni train/test assicurano che nessun near-homolog appaia in entrambi i set.

3. Architettura e Addestramento SAE:

   • Architettura: Viene impiegato un Matryoshka BatchTopK SAE. Questo combina:
     • Matryoshka: Dimensioni del prefisso annidate che permettono la ricostruzione a molteplici granularità.
     • BatchTopK: Impone la sparsità mantenendo solo i top $K$ pre-attivazioni attraverso un batch, sostituendo le penalità L1.
   • Iperparametri: Dimensione del dizionario $m=12.288$ (espansione di 16x), $K=80$ latenti attivi, addestrato per 20.000 step.
   • Performance: Gli SAE addestrati spiegano il 96,9% della varianza delle attivazioni tenute fuori dal training (held-out) con una sparsità $L_0$ media di 79,9.

4. Protocollo di Valutazione:

   • Probing: I vettori latenti SAE ($z$) mediati (mean-pooled) e le attivazioni grezze ($x$) vengono utilizzati per addestrare probe di regressione logistica con regolarizzazione L2.
   • Metriche: L'AUROC è valutato sotto due regimi: split stratificati casuali e split raggruppati per omologia.
   • Feature Scoring: L'AUROC univariata è calcolata per ogni feature rispetto alle etichette di classe. La significatività è determinata tramite test di Mann–Whitney U con correzione FDR di Benjamini–Hochberg ($q < 0,05$).

Risultati Chiave

1. I Latenti SAE Superano le Attivazioni Grezze:

   • A RFD3 Blocco 12, le caratteristiche codificate da SAE superano significativamente le attivazioni grezze nel setting di omologia raggruppata, raggiungendo un'AUROC di 0,817 rispetto a 0,763 per le feature grezze (un miglioramento di +0,054).
   • Questo miglioramento è specifico per il Blocco 12; i blocchi precedenti (6, 8) mostrano guadagni neutri o negativi dai SAE.
   • Per RF3 (dataset ToxinPred3), le feature del Blocco 12 SAE mostrano anche un miglioramento di +0,041 rispetto alle feature grezze negli split raggruppati.

2. Riduzione della Memorizzazione dell'Omologia:

   • Il divario tra l'AUROC dello split casuale e quello dello split raggruppato per omologia (un proxy per la memorizzazione della famiglia) è minimo per le feature SAE di RFD3 Blocco 12 (gap di 0,13) rispetto al Blocco 6 (gap di 0,22). Ciò suggerisce che i SAE del Blocco 12 catturano segnali strutturali discriminanti per classe e generalizzabili alla famiglia, piuttosto che memorizzare le famiglie di fold.

3. Identificazione di Feature di Pericolo Monosemantiche:

   • Lo studio identifica singole feature SAE che si attivano selettivamente su design pericolosi.
   • Top Feature: La Feature 639 in RFD3 Blocco 12 raggiunge un'AUROC di 0,819 ($q \approx 10^{-13}$) per il rilevamento di neurotossine di vipera.
   • Localizzazione: Queste feature mostrano una localizzazione per residuo. Per la Feature 639, i primi 3 residui attivanti (99, 104, 109) si raggruppano all'interno di una singola alfa-elica della struttura di Ammodytoxin A. Altre feature si localizzano similmente in regioni specifiche di proteine di evasione immunitaria di Monkeypox, proteine inattivanti i ribosomi e tossine salivari di zanzara.

4. Performance Competitive:

   • Il miglior probe VFUSE (RFD3 Blocco 12 SAE, split casuale, AUROC 0,877) è competitivo con i classificatori state-of-the-art basati solo su sequenza (ad esempio, DTVF a 0,92), nonostante utilizzi un semplice probe logistico e un dataset più piccolo. Fondamentalmente, VFUSE fornisce razionali per residuo, che i modelli basati solo sulla sequenza non possiedono.

Significato e Rivendicazioni

Il paper sostiene di essere il primo SAE addestrato su un modello di diffusione all-atom e il primo audit di pericolosità a livello di feature di un modello di design proteico.

• Insight Meccanicistico: Il lavoro dimostra che l'apprendimento del dizionario a profondità specifiche (Blocco 12) scompone concetti strutturali polisemantici in direzioni monosemantiche che sono più adatte al rilevamento di pericoli con bassa capacità rispetto alle rappresentazioni grezze del modello.
• Sicurezza e Interpretabilità: VFUSE stabilisce un percorso per il monitoraggio in tempo reale e l'audit interpretabile dei modelli generativi di proteine. Identificando residui e motivi strutturali specifici associati alla virulenza, il metodo va oltre lo screening "black box" verso controlli di sicurezza spiegabili.
• Direzioni Future: Sebbene gli autori notino limitazioni (dataset piccolo, mancanza di successo nel controllo/steering nei test preliminari), il lavoro apre la strada a un design proteico sicuro e interpretabile, consentendo il rilevamento di caratteristiche pericolose a livello di feature piuttosto che solo a livello di sequenza.