From Words to Amino Acids: Does the Curse of Depth Persist?

1. Enunciato del Problema

Le recenti analisi sui Modelli Linguistici su Larga Scala (LLM) hanno identificato la "Maledizione della Profondità", un fenomeno per cui gli strati più profondi nei transformer autoregressivi contribuiscono minimamente alle previsioni finali, con la maggior parte del calcolo rilevante per il compito che avviene negli strati iniziali o intermedi. Gli strati successivi servono principalmente a rifinire in modo incrementale la distribuzione di output piuttosto che generare nuove informazioni.

La domanda centrale di questo lavoro è se questa inefficienza di profondità persista nei Modelli Linguistici per le Proteine (PLM). A differenza degli LLM, i PLM utilizzano spesso:

• Obiettivi non autoregressivi: Come la Modellazione Linguistica Mascherata (MLM) e la Diffusione Discreta.
• Input multimodali: Che incorporano struttura e funzione proteica insieme alla sequenza.
• Paradigmi di scalatura diversi: Compresi architetture Mixture-of-Experts (MoE).

Gli autori indagano se la "Maledizione della Profondità" sia una proprietà universale dei modelli basati su Transformer profondi o se i vincoli unici della modellazione proteica (struttura biologica, vincoli evolutivi) alterino il modo in cui la profondità viene utilizzata.

2. Metodologia

Gli autori hanno condotto un'analisi sistematica della profondità su 7 famiglie di PLM comprendenti 26 varianti di modello, coprendo un'ampia gamma di dimensioni (da 8M a 15B parametri) e obiettivi di addestramento (Autoregressivo, Mascherato, Diffusione e Multimodale).

Hanno impiegato un framework unificato che combina tre tipi di misurazioni:

A. Analisi Interventale (Salto degli Strati)

Adattando il metodo "skiplayer" dalla ricerca sugli LLM, gli autori hanno misurato l'effetto propagato del salto di strati specifici:

• Modelli autoregressivi: Hanno saltato strati per i token iniziali e misurato l'impatto sulle previsioni dei token futuri.
• Modelli Mascherati/Diffusione: Hanno saltato strati su un sottoinsieme di token mascherati e misurato l'impatto sui restanti token mascherati non intervenuti.
• Metrica: La variazione relativa negli stati nascosti a valle e nelle norme di output.

B. Analisi della Similarità delle Rappresentazioni

• Similarità Strato-Strato: Calcolata la CKA Lineare (Centered Kernel Alignment) tra le rappresentazioni di diversi strati per quantificare quanto cambia l'informazione mentre si propaga attraverso la rete.
• Similarità Cross-Modale: Per i modelli multimodali (ESM3, DPLM2), hanno calcolato la Similarità del Coseno tra rappresentazioni solo-sequenza e solo-struttura negli strati corrispondenti per comprendere come le modalità si allineano e divergono.

C. Evoluzione della Distribuzione di Output

Utilizzando LogitLens e TunedLens, gli autori hanno estratto la distribuzione di previsione dei token implicata a ogni strato.

• Metriche: Hanno misurato la Divergenza KL tra le distribuzioni degli strati intermedi e l'output finale, e la Sovrapposizione Top-1 (accordo sul token più probabile).

D. Valutazione dell'Utilità a Valle

Per collegare le misurazioni intrinseche alle prestazioni pratiche, hanno valutato le prestazioni strato per strato su ProteinGym, un benchmark su larga scala per la previsione dell'effetto delle mutazioni zero-shot (217 saggi di scansione mutazionale profonda). Hanno misurato la correlazione di rango di Spearman tra le previsioni fatte dagli strati di uscita anticipata e i dati sperimentali.

3. Contributi Chiave

1. Primo Studio Sistematico sulla Profondità nei PLM: Il paper fornisce la prima analisi completa dell'uso della profondità attraverso diverse famiglie di PLM, obiettivi e modalità.
2. Framework Unificato: Ha sviluppato una metodologia coerente per confrontare modelli autoregressivi, mascherati, a diffusione e multimodali, dimostrando che l'inefficienza di profondità non è limitata alla previsione del token successivo.
3. Scoperta di un'Inefficienza di Profondità Generalizzata: Ha dimostrato che la "Maledizione della Profondità" è una caratteristica comune dei moderni PLM, dove un sottoinsieme di strati (spesso intermedi) cattura la maggior parte dei segnali rilevanti a valle, mentre gli strati successivi offrono rendimenti decrescenti.
4. Approfondimenti Specifici per Modalità: Ha rivelato profili di profondità distinti nei modelli multimodali (in particolare ESM3), spiegando il loro comportamento attraverso le dinamiche di allineamento cross-modale.

4. Risultati Chiave

A. La Maledizione della Profondità Persiste nei PLM

In tutte le famiglie di modelli (ESM2, DPLM, ProGen, ecc.), gli autori hanno osservato un pattern coerente:

• Strati Iniziali-Medi: Saltare questi strati causa una significativa interruzione dei calcoli a valle e riduce drasticamente le prestazioni a valle.
• Strati Tardi: Saltare una grande frazione degli strati finali (ad esempio, gli ultimi 30-50%) risulta in minime variazioni delle rappresentazioni a valle e in soli cali marginali delle prestazioni su ProteinGym.
• Effetto di Scalatura: All'aumentare delle dimensioni del modello, questa inefficienza diventa più pronunciata. Nei modelli grandi (ad esempio ESM2-15B), il "ginocchio" in cui le prestazioni si saturano si verifica prima nella profondità relativa.

B. Rifinitura della Distribuzione di Output

• Analisi TunedLens: La divergenza KL tra le distribuzioni degli strati intermedi e l'output finale diminuisce costantemente con la profondità.
• Stabilità: La previsione Top-1 spesso si stabilizza ben prima dell'ultimo strato. Gli strati successivi rifiniscono principalmente la confidenza (massa di probabilità) della previsione piuttosto che cambiare il token previsto stesso.

C. Modelli Multimodali e a Diffusione

• Generalizzazione: Il pattern di inefficienza di profondità vale per i Modelli a Diffusione (DPLM, DPLM2) e i Modelli Multimodali (ESM3, DPLM2) quando elaborano input solo-struttura o sequenza-struttura congiunti.
• Anomalia ESM3: ESM3 esibisce un profilo di profondità distinto. A differenza di altri modelli dove gli strati iniziali sono critici, ESM3 mostra effetti propagati deboli negli strati iniziali e contributi più forti che emergono negli strati medi-tardi.
  • Spiegazione: L'analisi della similarità cross-modale ha rivelato che i primi ~60% degli strati di ESM3 sono dedicati ad allineare le rappresentazioni di sequenza e struttura in uno spazio condiviso. Una volta allineate, gli strati successivi eseguono il calcolo specifico del compito.

D. Architetture Sparse (MoE)

• ProGen3: Come modello Mixture-of-Experts sparso, ProGen3 ha mostrato una "maledizione della profondità" meno pronunciata rispetto ai modelli densi. La separazione tra strati ad alto effetto e strati a basso effetto non si è affinata in modo significativo con la scala, suggerendo che la sparsità potrebbe mitigare l'inefficienza di profondità.

5. Significato e Implicazioni

• Efficienza Architettonica: I risultati suggeriscono che i moderni PLM profondi sono computazionalmente inefficienti. Una porzione significativa del costo di inferenza (FLOPs) è spesa su strati che contribuiscono poco all'output finale.
• Direzioni Future: Il paper motiva lo sviluppo di:
  • Calcolo Adattivo: Meccanismi di uscita anticipata o politiche di profondità dinamica che saltano gli strati ridondanti durante l'inferenza.
  • Architetture in Crescita/Looping: Modelli che espandono la profondità durante l'addestramento o riutilizzano gli strati (looping) per migliorare l'utilizzo.
  • Modelli Sparsi: Il successo di ProGen3 suggerisce che le architetture MoE potrebbero essere un percorso più efficiente per la scalatura dei modelli proteici.
• Progettazione Multimodale: Il comportamento distinto di ESM3 evidenzia che, in contesti multimodali, la profondità iniziale può essere consumata dall'allineamento delle modalità piuttosto che dal calcolo del compito, suggerendo che i modelli futuri potrebbero beneficiare di moduli specializzati di "allineamento" seguiti da stack di "calcolo" specifici per il compito.

In conclusione, la "Maledizione della Profondità" non è unica degli LLM autoregressivi ma è una caratteristica fondamentale dei moderni modelli proteici basati su Transformer profondi, indipendentemente dal loro obiettivo di addestramento o modalità. Questa intuizione richiede un cambiamento di paradigma verso architetture di IA proteica più efficienti in termini di profondità.