Riepilogo Tecnico: SAE mediati per turno per la scoperta di feature e l'attribuzione a lungo contesto

1. Definizione del problema

Gli Autoencoder Sparsi (SAE) sono emersi come uno strumento primario per estrarre feature interpretabili dai modelli linguistici, facilitando l'analisi dei circuiti, il controllo comportamentale (steering) e il monitoraggio della sicurezza. Tuttavia, le architetture SAE standard operano sulle attivazioni di singoli token. Questa granularità per-token crea due colli di bottiglia significativi per l'interpretabilità su larga scala:

1. Limitazioni di scalabilità: Il numero di feature attive scala linearmente con la lunghezza del contesto. In scenari a lungo contesto o conversazioni multi-turno, il volume di dati risultante rende l'analisi a valle (come i grafi di attribuzione) ingombrante e difficile da ispezionare.
2. Perdita di contesto di alto livello: Molti obiettivi di interpretabilità riguardano la comprensione del perché un modello abbia prodotto una specifica risposta, richiedendo l'identificazione di proprietà di alto livello (es. argomento, persona, tratti di sicurezza) che sono distribuite su molti token. Gli SAE per-token standard spesso apprendono detector di basso livello per parole specifiche o pattern sintattici. Poiché l'obiettivo di addestramento ricostruisce ogni posizione indipendentemente, non vi è alcuna pressione affinché il modello separi queste caratteristiche di alto livello, estese su interi intervalli (span), dal rumore a livello di token.

2. Metodologia

2.1 SAE mediati per turno (Turn-Averaged SAEs)

La proposta centrale è quella di addestrare gli SAE sulla media dello stato nascosto attraverso un intervallo di token (specificamente, un singolo turno dell'Umano o dell'Assistente) piuttosto che sulle attivazioni dei singoli token.

• Trasformazione dell'input: Dato un turno $T$ con token $x_t$, l'input è la media degli stati nascosti $\bar{x} = \frac{1}{|T|} \sum_{t \in T} x_t$.
• Architettura: L'architettura SAE rimane identica alle implementazioni standard (usando BatchTopK per la sparsità). Poiché la media di vettori in $\mathbb{R}^{d_{model}}$ mantiene la stessa dimensionalità dei singoli vettori, non sono richiesti cambiamenti strutturali.
• Intuizione: La media agisce come un filtro passa-basso. Leviga i dettagli ad alta frequenza e specifici del token (come la scelta di parole specifiche o marcatori sintattici) preservando al contempo i componenti a bassa frequenza che corrispondono a concetti ampi come argomento, stile e tratti comportamentali.

2.2 Architetture Ibride

Per affrontare il limite secondo cui la media scarta il dettaglio per-token, gli autori propongono due approcci ibridi per catturare sia la struttura di alto livello che i residui fini:

1. SAE Decoppiato (Decoupled SAE): Addestra due SAE indipendenti. Uno ricostruisce la media del turno ($\bar{x}$), e l'altro ricostruisce il residuo ($x_t - \bar{x}$). Sebbene questo separi i segnali, i modelli sono causalmente disconnessi, complicando l'analisi di attribuzione.
2. SAE Annidato (Nested SAE, stile Matryoshka): Addestra un singolo SAE con un encoder e un decoder condivisi ma una funzione di perdita specializzata. Lo spazio delle feature è diviso in una partizione interna $[0:h]$ e una partizione esterna $[h:d_{sae}]$.
   • Funzione di perdita: $L = \sum_t \| \text{dec}(\text{enc}(x_t)) - x_t \|^2 + \alpha \cdot \| \text{dec}_h(\text{enc}(\bar{x})[0:h]) - \bar{x} \|^2$.
   • La partizione interna è costretta a ricostruire la media del turno, mentre l'insieme completo delle feature ricostruisce le attivazioni per-token. Ciò consente a un singolo modello di apprendere simultaneamente sia le caratteristiche di alto livello del turno che i dettagli a livello di token.

2.3 Adattamento dei Grafi di Attribuzione

Gli autori adattano i grafi di attribuzione al livello del turno. Invece di nodi che rappresentano le feature in specifiche posizioni di token, i nodi rappresentano le feature attive su un particolare turno e layer.

• Pesi degli archi (Edge Weights): Calcolati mediando l'attribuzione del gradiente attraverso la sorgente del turno e normalizzando per la lunghezza del turno ($1/|T_A|$ e $1/|T_B|$) per garantire che i pesi siano agnostici rispetto alla lunghezza del turno.
• Validazione dell'intervento: I pesi degli archi sono validati perturbando le feature sorgente lungo le loro direzioni del decoder attraverso tutti i token di un turno e misurando il cambiamento nelle pre-attivazioni della feature target.

3. Contributi Chiave

1. SAE mediati per turno: Dimostrato che l'addestramento sugli stati nascosti medi produce feature che catturano le caratteristiche di alto livello del turno (argomento, stile, tratti di sicurezza) in modo più completo rispetto alle feature per-token.
2. Architettura SAE Annidata: Introdotta un'architettura a singolo modello che apprende congiuntamente le feature mediate per turno e quelle per-token utilizzando una perdita tipo Matryoshka, raggiungendo un'alta qualità di ricostruzione per entrambi i componenti del segnale.
3. Attribuzione Scalabile: Adattati i grafi di attribuzione al livello del turno, consentendo un'analisi compatta delle conversazioni multi-turno. Gli autori hanno validato questi grafi attraverso esperimenti di intervento e metriche di sufficienza (completezza e sostituzione).

4. Risultati

4.1 Qualità delle Feature e Ricostruzione

• Ricostruzione: Il Nested SAE raggiunge un nMSE (Normalized Mean Squared Error) per-token di 0,167, che corrisponde strettamente al puro SAE per-token (0,162), pur raggiungendo anche un nMSE mediato per turno di 0,116. Al contrario, un SAE per-token standard performa scarsamente sugli input mediati (nMSE 0,355), indicando che la media è fuori distribuzione (out-of-distribution) per i modelli addestrati su token.
• Scoperta Qualitativa: In esempi riguardanti errori matematici, stili tecnici entusiastici e conformità normativa, le feature mediate per turno hanno identificato con successo proprietà comportamentali di alto livello (es. "risposte errate", "stile tecnico entusiasta") che le feature per-token avevano mancato o catturato solo parzialmente.
• Generalizzazione a Lungo Contesto: Gli SAE mediati per turno addestrati su turni di chat (lunghezza media ~150 token) si sono generalizzati con successo a documenti 150 volte più lunghi (24k–32k token) da LongBench v2, identificando concetti di alto livello come "macroeconomia" e "grafica computerizzata" che le feature per-token non riuscivano a rilevare.

4.2 Valutazione Quantitativa

Utilizzando un framework "LLM-as-a-judge" per confrontare le descrizioni delle feature con i riassunti strutturati dei turni:

• Discriminazione (matching a 10 vie): Le feature per-token (max-pooled) hanno superato le feature mediate per turno (95,0% vs 73,9%), poiché queste ultime si affidano a dettagli specifici del livello di token per distinguere i testi.
• Copertura (classifica pairwise e a 5 vie): Le feature mediate per turno hanno superato significativamente le feature per-token (vincendo il 87,9% dei confronti pairwise). Fornivano descrizioni più complete delle caratteristiche di alto livello di un turno.
• Metrica di Embedding: Le feature mediate per turno hanno ottenuto la massima similarità del coseno (0,663) rispetto ai campi dei riassunti strutturati, confermando la loro superiore copertura del contenuto semantico.

4.3 Analisi di Attribuzione

• Compattezza del Grafo: I grafi mediati per turno hanno ridotto il conteggio dei nodi da circa 128K (per-token) a circa 5K per una conversazione di 10 turni, rendendo l'ispezione fattibile.
• Casi di Studio:
  • Attribuzione del System Prompt: I grafi mediati per turno hanno correttamente identificato quali paragrafi del system prompt influenzavano specifiche query dell'utente, mostrando shift di attribuzione comparabili ai grafi per-token.
  • Tracciamento della Degenerazione: In una conversazione di 14 turni in cui il modello è degenerato in un testo ripetitivo, il grafo mediato per turno ha tracciato una catena causale coerente: "Disclaimer di sicurezza AI" $\to$ "Armi nucleari" $\to$ "Rilevamento prompt avversari" $\to$ "Output degenerato". I grafi per-token hanno fallito nel trovare la feature di rilevamento avversario, tracciando invece verso concetti non correlati come "composti chimici".
• Intervento e Sufficienza: I pesi degli archi hanno mostrato una correlazione significativa con gli effetti dell'intervento. Il Nested SAE ha ottenuto i punteggi di sufficienza di "completezza" e "sostituzione" più elevati, superando sia il puro per-token che le aggregazioni post-hoc di feature per-token.

5. Significato e Rivendicazioni

Il paper sostiene che gli SAE mediati per turno rendono pratiche le tecniche di interpretabilità a lungo contesto. Spostando l'unità di analisi dai token agli intervalli (span), il metodo:

• Produce feature che descrivono le caratteristiche di alto livello (argomento, persona, sicurezza) in modo più completo rispetto alle feature a livello di token.
• Consente la costruzione di grafi di attribuzione compatti e interpretabili per conversazioni multi-turno che sarebbero altrimenti troppo grandi da analizzare.
• Dimostra che modificare il segnale di addestramento (tramite la media) può costringere gli SAE a imparare feature a una granularità desiderata (concetti di alto livello) indipendentemente dal numero di token.
• Mostra che queste feature si generalizzano a intervalli molto più lunghi di quelli visti durante l'addestramento, suggerendo robustezza per l'analisi di modelli frontier.

Gli autori segnalano i limiti, tra cui la scala attuale degli esperimenti (modello 7B, corpus LMSYS-Chat-1M) e la dipendenza da LLM per le metriche di valutazione, identificando l'applicazione ai modelli frontier e la creazione di benchmark con verità fondamentale (ground-truth) come necessari lavori futuri.