Sintesi Tecnica: La struttura del task inverte la codifica dello stato per livelli nei modelli di sequenza

Definizione del Problema

L'interpretabilità meccanicistica dei modelli di sequenza tratta spesso le codifiche dello stato per livelli come tratti architettonici fissi. L'assunto prevalente è che i modelli ricorrenti (ad es. LSTM, GRU) e i modelli State-Space (SSM come Mamba) concentrino lo stato leggibile negli strati tardivi, mentre i modelli basati sull'attenzione (Transformer) distribuiscono l'accumulo dello stato attraverso la profondità. Questo articolo sfida la visione secondo cui queste strategie di codifica siano invarianti architettoniche stabili. Esso postula che la stessa architettura inverta il proprio profilo di codifica dello stato per livelli quando la struttura computazionale del task cambia. Inoltre, la letteratura spesso confonde due potenziali driver di questi pattern: la struttura algebrica (specificamente la commutatività) rispetto alla struttura computazionale (aggiornamenti di prefisso vs. tracciamento di uno stack). Gli autori mirano a disaggregare questi fattori e a determinare se la codifica dello stato sia una proprietà dell'architettura da sola o un'interazione tra architettura e task.

Metodologia

Lo studio impiega una suite sperimentale controllata attraverso cinque famiglie di architetture: Transformer, Mamba, Mamba-2, LSTM e GRU.

Task:
Gli autori utilizzano tre task formali di tracciamento dello stato e un benchmark semi-reale:

1. Parity (Parità): Un task di aggiornamento del prefisso a singolo bit, commutativo.
2. Composizione di permutazioni $S_3$: Un task di aggiornamento del prefisso a 6 classi, non commutativo. Questo funge da discriminante critico: se la commutatività algebrica guida il profilo di codifica, $S_3$ dovrebbe comportarsi come Dyck; se la struttura computazionale guida il profilo, $S_3$ dovrebbe comportarsi come Parity.
3. Dyck-$k$: Un task di tracciamento di uno stack a profondità limitata (corrispondenza di parentesi) con una profondità massima di 10.
4. Code-Depth: Un task di profondità di annidamento a livello di token derivato da codice Python in CodeSearchNet, che funge da benchmark semi-reale di validazione.

Configurazione Sperimentale:

• Task Formali: Gli esperimenti utilizzano modelli a 4 strati con $d_{model}=128$. Le architetture sono parametricamente accoppiate (circa 830K parametri) ove possibile. L'addestramento avviene con lunghezza di sequenza $L=40$, con valutazione a $L \in \{40, 60, 80, 100, 120\}$.
• Scala Pre-addestrata: Il fine-tuning viene eseguito su Mamba-130M (24 strati) e Pythia-160M/410M (12/24 strati).
• Misurazioni:
  • Linear Probing: Addestramento di classificatori di regressione logistica su stati nascosti congelati per misurare la leggibilità dello stato in ogni strato.
  • Interventi Causali: Zero-ablation (proiezione fuori dalle direzioni del probe) e subspace ablation (rimozione delle direzioni principali PCA o singolari) per testare la necessità funzionale.
  • Generalizzazione OOD: Applicazione di ablazioni addestrate a $L=40$ su lunghezze non viste ($L=80, 100, 120$).
  • Attribuzione dei Componenti: Ablazione del componente Conv1D di Mamba sotto un backend puramente PyTorch corrispondente per isolare il suo contributo attraverso i task.
  • Activation Patching: In Mamba pre-addestrato, l'applicazione di patch di attivazioni di posizione intermedia da input puliti in input corrotti per misurare il recupero del logit.

Contributi Chiave

1. Inversione Dipendente dal Task: Il paper dimostra che il pattern "concentrazione tardiva" vs. "costruzione distribuita" non è un invariante architettonico. Esso si inverte a seconda del task:
   • Task di Aggiornamento del Prefisso (Parity, $S_3$): I modelli ricorrenti e gli SSM concentrano lo stato leggibile nello strato finale, mentre i Transformer lo costruiscono gradualmente attraverso gli strati.
   • Tracciamento dello Stack Limitato (Dyck): Il pattern si inverte. I Transformer raggiungono una leggibilità quasi perfetta immediatamente (strati iniziali), mentre i modelli ricorrenti e gli SSM costruiscono la rappresentazione progressivamente.
2. Disaggregazione di Algebra e Computazione: Introducendo il task non commutativo $S_3$, gli autori mostrano che il profilo di codifica segue la struttura computazionale (aggiornamento del prefisso vs. stack) piuttosto che la struttura algebrica (commutatività). $S_3$ si raggruppa con Parity in tutte le architetture, smentendo l'ipotesi che la commutatività sia il driver primario del profilo di codifica.
3. Disaccoppiamento tra Leggibilità e Causalità: Lo studio rivela che lo strato in cui lo stato è più linearmente leggibile (miglior strato del probe) non coincide sempre con lo strato che è causalmente necessario per il task. Questo disaccoppiamento è particolarmente pronunciato nei modelli pre-addestrati, dove le rappresentazioni sono più ridondanti e distribuite.
4. Dipendenza dal Task a Livello di Componente: Anche all'interno di una singola architettura (Mamba), l'importanza di componenti specifici varia a seconda del task. Ad esempio, il componente Conv1D è critico per $S_3$ e Parity ma non ha un effetto significativo sulle prestazioni di Dyck, sfidando la nozione di un unico "meccanismo" per un'architettura.

Risultati Chiave

• Livello Comportamentale: I baseline ricorrenti generalizzano perfettamente su Parity e $S_3$ e superano SSM e Transformer su Dyck. L'ordine delle prestazioni di $S_3$ (Ricorrente > SSM > Transformer) rispecchia quello di Parity, non di Dyck.
• Profili di Probing:
  • Parity/$S_3$: Mamba e LSTM mostrano un salto netto nell'accuratezza del probe nell'ultimo strato ($L_3$ in modelli a 4 strati). I Transformer mostrano un aumento graduale da $L_1$ a $L_3$.
  • Dyck: I Transformer raggiungono un'alta accuratezza a $L_0$ o $L_1$. Mamba e LSTM mostrano un aumento graduale, raggiungendo la saturazione solo in strati successivi.
  • Conferma di $S_3$: $S_3$ segue il profilo di Parity (concentrazione tardiva per RNN/SSM, graduale per Transformer), confermando l'ipotesi della struttura computazionale.
• Interventi Causali:
  • Nei modelli formali a 4 strati, l'ablazione della migliore direzione del probe causa cali significativi di accuratezza, e queste direzioni rimangono necessarie a lunghezze di sequenza maggiori (OOD).
  • Nei Modelli Pre-addestrati, si verifica una scissione:
    • Pythia (Dyck): Esiste un forte bottleneck negli strati medi (l'ablazione a $L_6$–$L_7$ riduce l'accuratezza dell'81% a 160M). Il miglior strato del probe ($L_{11}$) non è quello più criticamente causale.
    • Mamba (Parity/Dyck/$S_3$): L'ultimo strato è altamente leggibile, ma l'ablazione di una singola direzione causa cali trascurabili (<0.2%). Tuttavia, l'activation patching di posizione intermedia nell'ultimo strato recupera circa il 97–98% del gap del logit, indicando che lo stato è sufficiente ma distribuito su molte direzioni, non catturato da un singolo probe lineare.
• Attribuzione Conv1D: Rimuovere Conv1D da Mamba causa un calo significativo su $S_3$ (-12.64 pp) e Parity (-6.83 pp), ma nessun calo significativo su Dyck (+4.04 pp, non significativo). Ciò indica che l'importanza del componente è dipendente dal task.

Significato e Rivendicazioni

L'articolo sostiene che chiedere "come l'architettura X codifica lo stato?" è incompleto senza specificare la struttura computazionale del task. Gli autori affermano che:

1. Le firme meccanicistiche sono proprietà congiunte: Le strategie di codifica dello stato non sono tratti architettonici fissi, ma emergono dall'interazione tra i bias induttivi del modello e i requisiti computazionali del task (aggiornamento del prefisso vs. stack).
2. Cautela metodologica: Estrapolare i risultati meccanicistici da un regime di task (es. aggiornamenti di prefisso) a un altro (es. tracciamento di uno stack a profondità limitata) rischia di scambiare firme specifiche del task per invarianti architettoniche.
3. Leggibilità vs. Causalità: Il probing identifica dove l'informazione è linearmente disponibile, ma questo non mappa sempre dove la computazione è in un bottleneck, specialmente nei grandi modelli pre-addestrati dove la ridondanza è alta.
4. Algebra vs. Computazione: Nel contesto di questi task formali, la struttura computazionale (aggiornamento sequenziale del prefisso vs. stack gerarchico) è un predittore migliore dei profili di codifica per strati rispetto alle proprietà algebriche come la commutatività.

Gli autori concludono che le loro scoperte valgono per i task formali, le scale pre-addestrate (fino a 410M di parametri) e i benchmark di codice semi-reali, suggerendo una robusta interazione tra architettura e struttura del task nella modellazione delle sequenze.