Sintesi Tecnica: I contatti proteici sono già presenti nell'attenzione

Definizione del Problema

I modelli linguistici per proteine (PLM), addestrati su vasti corpora di sequenze non etichettate, apprendono rappresentazioni ricche della struttura proteica. Una domanda centrale è se tale conoscenza strutturale possa essere estratta senza un addestramento supervisionato su etichette strutturali. Zhang et al. [2024] hanno proposto il Jacobiano Categoriale Standard (CJ) come soluzione: perturbando ogni residuo con ogni aminoacido alternativo e misurando la risposta nei logit del modello di masking-language-model (MLM), è possibile ricostruire una mappa di contatto. Sebbene il CJ raggiunga un'alta precisione (P@L ≈ 0,80 su determinati benchmark), è computazionalmente costoso, richiedendo circa $19L$ passaggi in avanti (forward pass) per una proteina di lunghezza $L$. Inoltre, il CJ è strutturalmente legato ad architetture con un head categoriale per posizione (bidirectional masked-LMs), limitandone l'applicabilità a modelli encoder-decoder o causali.

Valutazioni recenti hanno inoltre sollevato preoccupazioni riguardanti la fuga di dati (data leakage), dove i corpora di pretraining possono sovrapporsi ai set di valutazione, gonfiando le metriche di performance. È necessario un metodo che sia computazionalmente efficiente, applicabile a diverse architetture e robusto rispetto alla fuga dal pretraining.

Metodologia

Gli autori propongono che il segnale di contatto ricostruito dal Jacobiano Categoriale sia già esplicitamente presente nelle mappe di attenzione del modello, concentrato in un piccolo sottoinsieme identificabile di teste di attenzione.

1. Fusione dell'Attenzione Top-K

Invece di perturbare gli input, il metodo proposto estrae i segnali di contatto direttamente dai meccanismi di attenzione del PLM:

• Estrazione: Per ogni testa di attenzione, la mappa di attenzione viene simmetrizzata, priva di diagonale (zero-diagonalized) e corretta tramite l'Average Product Correction (APC).
• Selezione: Le teste vengono classificate in base alla loro capacità di predire i contatti (misurata tramite la precisione a lungo raggio top-L/2) su un piccolo set di selezione etichettato e disgiunto (fino a 10 proteine).
• Fusione: Le top-$K$ teste selezionate vengono combinate utilizzando una semplice media aritmetica non pesata. Ciò richiede un solo passaggio in avanti per proteina durante l'inferenza.
• Efficienza delle Etichette: La selezione delle top-$K$ teste è l'unico passaggio supervisionato, richiedendo un budget minimo di proteine etichettate (10 per la selezione, 50 per un controllo a budget equivalente rispetto alla regressione supervisionata).

2. Representation-CJ (Repr-CJ)

Per confrontarsi con il CJ su architetture prive di un head MLM (ad esempio, modelli encoder-decoder o causali), gli autori introducono il Representation-CJ. Questo generalizza la costruzione del Jacobiano sostituendo la risposta del logit con la risposta dello stato nascosto:
$$J_{repr}[i, a', j] = \| h_{i \to a'}[\ell, j] - h_{unpert}[\ell, j] \|_2$$
dove $h$ rappresenta lo stato nascosto alla posizione $j$ e al livello $\ell$ quando il residuo $i$ viene perturbato in l'aminoacido $a'$. Questo mantiene lo stesso costo computazionale ($19L + 1$ passaggi in avanti) del logit-CJ, ma si applica a qualsiasi architettura che esponga gli stati nascosti.

3. Protocollo di Valutazione

Lo studio utilizza una strategia di valutazione leakage-clean per garantire che i risultati non siano artefatti della sovrapposizione del pretraining:

• CAMEO-PTA25: Un set di valutazione primario costituito da catene proteiche depositate dopo il cutoff di pretraining dei modelli (post-2022-09), filtrate contro l'attuale UniRef50 con identità al 50%.
• Selezione Disgiunta: La selezione delle teste e la valutazione sono eseguite su sottoinsiemi di dati disgiunti per prevenire la fuga di informazioni.
• Baseline: Il metodo viene confrontato con lo Standard Categorical Jacobian (logit-CJ), la migliore singola testa di attenzione e le baseline di regressione logistica L1 supervisionata.

Contributi Chiave

1. Efficienza e Performance: Il metodo di fusione dell'attenzione Top-K recupera i segnali di contatto in un singolo passaggio in avanti, superando il Categorical Jacobian in termini di dati leakage-clean per ogni modello bidirezionale testato (famiglia ESM-2, ESM-1b, ProtT5-XL).
2. Validazione Leakage-Clean: Su uno split CAMEO-clean, il metodo di fusione supera significativamente il CJ (ad esempio, +9 punti percentuali su ESM-2-650M). Fondamentalmente, su questi dati puliti, il Categorical Jacobian non supera più la migliore singola testa di attenzione, suggerendo che la sua precedente superiorità in-distribution fosse in parte dovuta alla memorizzazione.
3. Generalizzazione Architetturale: Gli autori dimostrano che il segnale di contatto esiste nei modelli encoder-decoder (ProtT5-XL) tramite la fusione dell'attenzione e introducono Repr-CJ per estendere l'analisi basata sul Jacobiano a queste architetture.
4. Insight Meccanicistico (Diffusione): Il numero ottimale di teste ($K$) da fondere dipende dalla "diffusione" del segnale di contatto all'interno del modello.
   • Cluster Concentrati: Alcuni modelli (ad esempio, ESM-2-8M, AMPLIFY-350M) concentrano il segnale in una singola testa dominante; mediare più teste diluisce il segnale.
   • Cluster Diffusi: Modelli più grandi (ad esempio, ESM-2-650M, ProtT5-XL) distribuiscono il segnale su molte teste; la mediazione è necessaria per recuperare il consenso.
5. Limitazione dei Modelli Causali: Sia il metodo di fusione che il metodo CJ falliscono sui modelli linguistici causali (ProGen2-large e -xlarge), suggerendo che il pretraining bidirezionale sia necessario affinché i meccanismi di attenzione codifichino strutture di coppia simmetriche.

Risultati

• Gap di Performance: Sui dati leakage-clean, la fusione batte il CJ di circa 9 punti percentuali su ESM-2-650M. Entrambi i metodi subiscono un calo di 30–36 punti percentuali in precisione assoluta passando dal benchmark in-distribution (Zhang eval-200) a quello leakage-clean (CAMEO-PTA25), indicando che i precedenti risultati di predizione di contatti non supervisionata erano significativamente gonfiati dalla sovrapposizione del pretraining.
• Supervisionato vs Non Supervisionato: A un budget di etichette equivalente (50 proteine), una media non pesata delle top-$K$ teste performa in modo statisticamente indistinguibile da un fit di regressione logistica L1 supervisionata sulle stesse teste. Ciò implica che la selezione delle teste è il driver primario della performance, non la forma parametrica del readout.
• Economia Computazionale: La fusione richiede un singolo passaggio in avanti per proteina, mentre il CJ ne richiede $\sim 19L$. Il punto di pareggio per il costo computazionale è inferiore a 20 proteine di test, rendendo la fusione enormemente più efficiente per applicazioni su larga scala.
• Validazione Repr-CJ: Il Repr-CJ correla strettamente con il logit-CJ (Pearson $r > 0.90$) sulle architetture in cui entrambi sono definiti, validandolo come una generalizzazione fedele per i modelli non MLM.

Significato e Rivendicazioni

Il paper sostiene che il Categorical Jacobian non estrae nuove informazioni di contatto, ma piuttosto ricostruisce l'informazione di coppia strutturale che il modello già rappresenta nei suoi pattern di attenzione. La scoperta principale è che questo segnale è accessibile tramite una semplice mediazione senza parametri di un piccolo sottoinsieme di teste di attenzione, offrendo un'alternativa computazionalmente superiore ai metodi basati sulla perturbazione.

Gli autori sottolineano che la valutazione "leakage-clean" rivela un divario sostanziale tra la performance in-distribution e quella di vera generalizzazione, suggerendo che i futuri benchmark per la predizione di contatti non supervisionata debbano adottare come standard split consapevoli del pretraining. Il lavoro evidenzia inoltre che l'allocazione dell'informazione strutturale tra le teste di attenzione è una proprietà appresa del processo di pretraining piuttosto che un'inevitabilità architettonica, variando anche all'interno della stessa famiglia di modelli (ad esempio, la scala di ESM-2).

Infine, il collasso di entrambi i metodi sui modelli linguistici causali è presentato come prova che il pretraining bidirezionale è un fattore critico per apprendere strutture di coppia simmetriche nell'attenzione, sebbene gli autori notino che questa sia una scoperta associativa dato il scope limitato dei modelli causali testati.