Sintesi Tecnica: Tokenizzazione Proteica Adattiva (APT)

Problema

Gli approcci attuali alla tokenizzazione della struttura proteica si basano sulla località, dove i token vengono creati aggregando informazioni da residui spazialmente vicini lungo una sequenza proteica. Sebbene questi metodi raggiungano un'alta fedeltà di ricostruzione, gli autori identificano due criticità fondamentali:

1. Accumulo di Errore nella Generazione: I modelli generativi basati sulla tokenizzazione locale soffrono di propagazione dell'errore. Un singolo campionamento errato di un token può portare a deviazioni significative nella struttura predetta, causando prestazioni inferiori dei modelli discreti rispetto ai modelli di diffusione continui nelle metriche generative.
2. Scarsa Scalabilità: Il numero di token scala linearmente con la dimensione della proteina. Ciò rende la modellazione di grandi complessi proteici computazionalmente costosa e limita la capacità di comprimere efficacemente le proteine.
3. Limitazioni della Rappresentazione: I tokenizer locali spesso richiedono operazioni di mean-pooling per generare rappresentazioni a dimensione fissa per task a valle, il che può scartare informazioni globali critiche e scalare male con la dimensione della proteina.

Metodologia

Gli autori propongono la Tokenizzazione Proteica Adattiva (APT), un metodo di tokenizzazione globale in cui i successivi token contribuiscono con livelli crescenti di dettaglio a una rappresentazione globale, anziché corrispondere a specifici vicinati locali.

Architettura

APT è implementata come un autoencoder di diffusione con un collo di bottiglia discreto:

• Input: Una sequenza di coordinate $C_\alpha$ grezze (normalizzate a massa zero) di lunghezza $L$.
• Encoder: Un transformer con attenzione bidirezionale mappa le posizioni grezze in una sequenza latente $c \in \mathbb{R}^{L \times d}$.
• Discretizzazione: La sequenza latente è discretizzata utilizzando la Quantizzazione Scalare Finita (FSQ) con livelli $(8, 5, 5, 5)$, che risulta in una dimensione effettiva del codebook di 1.000.
• Meccanismo di Adattività: Per imporre l'adattività, il modello impiega un dropout nidificato. Durante l'addestramento, un limite superiore $U$ viene campionato uniformemente da $[1, \min(L, k_{max})]$. I token oltre questo limite vengono scartati. Ciò incoraggia il modello a collocare le informazioni globali critiche (componenti a bassa frequenza) nei primi token e i dettagli ad alta frequenza nei token successivi.
• Decoder: Un decoder di diffusione, addestrato utilizzando un obiettivo di flow-matching, ricostruisce le coordinate atomiche condizionate dai token discreti. Il decoder utilizza l'equivarianza stocastica, apprendendo le simmetrie piuttosto che imporle esplicitamente.
• Predizione della Dimensione: La dimensione della proteina viene regredita dal primo token tramite una perdita di cross-entropy, permettendo al modello di disaccoppiare la dimensione della proteina dalla lunghezza della sequenza di condizionamento.

Addestramento e Inferenza

• Addestramento: Il modello è addestrato su circa 473.000 predizioni sintetiche di AlphaFold2 provenienti dal database AFDB clusterizzato tramite Foldseek. L'addestramento prevede due fasi: (1) addestramento dell'autoencoder con rotazioni casuali per apprendere l'equivarianza stocastica, e (2) addestramento di un modello autoregressivo (stile GPT) sui token APT.
• Strategie di Inferenza:
  • Tail Dropout (Dropout della Coda): Durante la generazione, il modello può interrompere la generazione anticipatamente. Qualsiasi prefisso della sequenza di token è un segnale di condizionamento valido. La coda (dettagli ad alta frequenza) è delegata al decoder di diffusione.
  • Arresto Basato sull'Entropia: Gli autori introducono criteri di arresto basati sull'entropia del token (cutoff finito, entropia minima basata su spline) per determinare il punto ottimale per interrompere la generazione, bilanciando fedeltà della rappresentazione e capacità generativa.
  • Classifer Annealing (Annealing del Classificatore): Per mitigare gli artefatti causati dallo spostamento lontano dal vero manifold dei dati, gli autori utilizzano l'annealing del classificatore, interpolando tra i campi di guida condizionali e incondizionali durante la decodifica.

Contributi Chiave

1. Tokenizer Globale Adattivo: Introduzione di un tokenizer in cui i token rappresentano descrittori globali in una gerarchia da grossolano a fine, anziché vicinati locali.
2. Validazione delle Prestazioni: Dimostrazione che il modello eguaglia o supera i tokenizer locali allo stato dell'arte (es. DPLM2, ESM3, Kanzi) in compiti generativi, di apprendimento della rappresentazione e di ricostruzione.
3. Tecniche di Inferenza: Sviluppo del campionamento basato sull'entropia e dell'annealing del classificatore per gestire la complessità del campione e mitigare l'esposizione all'errore, portando a una migliore designabilità.
4. Applicazioni Zero-Shot: Dimostrazione del disaccoppiamento della dimensione proteica dal condizionamento, abilitando applicazioni zero-shot come il protein shrinking (riduzione del numero di residui mantenendo la struttura) e l'affinity maturation tramite scaling all'inferenza.

Risultati

• Ricostruzione: Sui set di test CATH, CAMEO e AFDB, APT raggiunge RMSD e score TM comparabili ai modelli di diffusione continua allo stato dell'arte. Notevolmente, utilizzando solo 32–64 token, si ottengono RMSD $< 2$ Å, confermando la comprimibilità delle strutture proteiche.
• Generazione: Un modello autoregressivo addestrato sui token APT raggiunge una designabilità di 0,871 (frazione di campioni con scRMSD $< 2$ Å), superando significativamente i modelli di diffusione discreta come DPLM2 (0,486) ed ESM3. Ciò avviene senza l'uso di tecniche di "best-of-N" sampling.
• Apprendimento della Rappresentazione: Nel compito di classificazione CATH, il probing non lineare (MLP) sui token APT supera il probing sulle rappresentazioni di ESM3 e DPLM2. Fondamentalmente, APT fornisce rappresentazioni globali a dimensione fissa senza mean-pooling, e anche rappresentazioni altamente compresse (16 token) superano modelli più grandi.
• Applicazioni:
  • Protein Shrinking: Ha generato con successo versioni più piccole di proteine (es. emoglobina, beta-barrel) preservando la struttura globale e locale condizionando il processo di diffusione su meno residui.
  • Affinity Maturation: Ha dimostrato la capacità di evolvere leganti deboli in leganti forti utilizzando la ricerca beam search guidata da funzioni di ricompensa (es. iPAE) nello spazio latente.

Significato

L'articolo sostiene che la Tokenizzazione Proteica Adattiva fornisca un approccio task-aware per scalare i compiti biologici. Passando dalla tokenizzazione locale a quella globale, il metodo risolve il compromesso tra fedeltà di ricostruzione e capacità generativa. La capacità di comprimere le proteine in vettori a dimensione fissa senza perdita di informazioni tramite mean-pooling offre una via pratica per modellare grandi complessi proteici e per eseguire lo scaling all'inferenza. Gli autori posizionano questo lavoro come un passo verso modelli generativi capaci di ragionare attraverso molteplici scale dimensionali, affrontando una sfida chiave nella bioingegneria di frontiera.

Nota: Gli autori riconoscono limitazioni, specificamente che i token APT sono efficaci per compiti globali ma meno adatti per compiti locali come lo scaffolding di motivi, e che l'attuale tokenizer non tiene conto delle catene laterali o della composizione amminoacidica, che sono critiche per certe funzioni.