Riepilogo Tecnico: Trasporto Ottimo Semantico per l'Abbinamento delle Caratteristiche degli Autoencoder Sparsi e la Compressione dei Circuiti

1. Definizione del Problema

Gli Autoencoder Sparsi (SAE) si sono affermati come strumento primario per l'interpretazione dei Modelli Linguistici di Grandi Dimensioni (LLM), decomponendo gli spazi latenti in caratteristiche interpretabili. Tuttavia, rimangono due sfide critiche nel scalare l'analisi degli SAE:

1. Abbinamento delle Caratteristiche tra Livelli: Identificare caratteristiche semanticamente simili tra diversi livelli di un LLM. I metodi esistenti si basano tipicamente sul confronto dei vettori del decoder. Questo approccio fallisce nell'abbinare caratteristiche tra livelli non consecutivi perché i vettori del decoder sono addestrati per ricostruire distribuzioni di attivazione distinte, rendendo il confronto diretto dei vettori inaccurato.
2. Compressione dei Circuiti di Caratteristiche: I circuiti di caratteristiche (grafici causali che collegano le caratteristiche tra i livelli) contengono spesso centinaia di nodi, rendendo l'ispezione manuale impraticabile. I metodi attuali di compressione si basano sul raggruppamento manuale di nodi semanticamente simili in "super-nodi", un processo che non scala.

Il problema centrale identificato dagli autori è la stima della distanza semantica tra caratteristiche che risiedono su varietà di attivazione diverse. Gli approcci esistenti trattano le caratteristiche come singoli vettori (ad esempio, pesi del decoder), il che scarta il ricco contesto di dove e quanto fortemente una caratteristica si attiva, portando a una perdita di informazioni e all'incapacità di gestire confronti multi-varietà.

2. Metodologia

Gli autori propongono un framework distribuzionale che riformula il confronto delle caratteristiche come un problema di Trasporto Ottimo (OT).

2.1 Rappresentazione delle Caratteristiche come Distribuzioni

Invece di rappresentare una caratteristica mediante un singolo vettore del decoder, il metodo rappresenta ciascuna caratteristica $i$ al livello $\ell$ come una distribuzione empirica ponderata sull'attivazione sugli stati nascosti in cui essa si attiva.

• Dato un corpus di posizioni di token $T$, il metodo seleziona le prime $K$ posizioni di token in cui la caratteristica $i$ si attiva con maggiore intensità.
• Ogni token selezionato $t$ viene assegnato un peso $w_{i,t}$ proporzionale alla sua magnitudine di attivazione.
• La caratteristica è rappresentata come una distribuzione $\hat{\mu}_{i,T} = \sum w_{i,t} \delta_{x_t}$, dove $x_t$ è lo stato nascosto al token $t$.

2.2 Spazio di Riferimento Condiviso

Per confrontare caratteristiche tra diversi livelli (che risiedono su varietà diverse), il metodo proietta queste distribuzioni specifiche per livello in uno spazio di riferimento condiviso $\mathcal{R}$.

• Per l'abbinamento tra livelli, lo spazio di riferimento è tipicamente lo spazio degli stati nascosti del livello target. Le attivazioni dal livello sorgente vengono proiettate sulle corrispondenti posizioni di token nel livello target.
• Per la compressione dei circuiti, lo spazio di riferimento è la concatenazione degli stati nascosti attraverso tutti i livelli.

2.3 Distanza Semantica tramite Distanza di Wasserstein

Una volta proiettate nello spazio condiviso, la distanza semantica tra due caratteristiche viene calcolata utilizzando la distanza di Wasserstein (Distanza del Movitore di Terra).

• La distanza $W_c(\hat{\mu}, \hat{\nu})$ misura il costo minimo per trasportare la massa di probabilità della distribuzione di attivazione di una caratteristica all'altra.
• Il costo di base $c(z, z')$ misura la spesa per abbinare due rappresentazioni di contesto nello spazio condiviso.
• Questo approccio cattura la struttura geometrica dei pattern di attivazione piuttosto che la semplice similarità vettoriale punto per punto.

2.4 Attività a Valle

• Abbinamento delle Caratteristiche: Per una caratteristica target, il metodo trova la caratteristica sorgente con la minima distanza di Wasserstein (vicino più prossimo).
• Compressione dei Circuiti: Il metodo tratta i nodi del circuito come punti nello spazio distribuzionale e applica il clustering (ad esempio, clustering agglomerativo) per raggruppare i nodi con piccole distanze di Wasserstein a coppie in super-nodi, minimizzando la distanza semantica all'interno del cluster.

3. Contributi Chiave e Garanzie Teoriche

Il documento fornisce tre principali garanzie teoriche per il framework proposto:

1. Invarianza alla Ridimensionamento delle Attivazioni: La distribuzione proiettata della caratteristica è invariante alla scalatura globale dei valori di attivazione. Poiché la rappresentazione si basa sulla classificazione relativa e sulla distribuzione proporzionale delle attivazioni (selezione Top-K e pesi normalizzati), spostamenti scalari arbitrari nelle magnitudini di attivazione non alterano la distanza semantica.
2. Stabilità sotto Perturbazioni: La distanza di Wasserstein tra le distribuzioni delle caratteristiche è stabile sotto perturbazioni lisce delle distribuzioni. La variazione della distanza è limitata dalle costanti di Lipschitz del costo di base e dalla distanza di Wasserstein tra le distribuzioni perturbate e originali. Questo garantisce robustezza contro il rumore nella selezione dei token o nel re-addestramento degli SAE.
3. Recupero Basato sul Margine: Gli autori stabiliscono le condizioni in cui l'abbinamento empirico del vicino più prossimo e il raggruppamento semantico recuperano perfettamente le corrispondenze vere (oracle). Se il "margine" (il divario tra la migliore corrispondenza e la seconda migliore) supera il doppio dell'errore di stima, la decisione empirica è garantita come corretta.

4. Risultati Sperimentali

Il metodo è stato valutato su SAE addestrati su GPT-2 e Gemma-2-2B, confrontandosi con baseline tra cui SAE Match (L2 del vettore del decoder), Feature Flow (coseno del vettore del decoder), Attribution Patching e selezione basata su LLM.

• Abbinamento delle Caratteristiche: Il metodo proposto ha superato tutte le baseline nella Valutazione LLM (valutazioni di similarità semantica) e nell'Accuratezza di Abbinamento attraverso configurazioni di livelli consecutivi, da medio a ultimo e da primo a ultimo. Crucialmente, ha mantenuto prestazioni competitive sulle metriche a valle (perdita di Entropia Incrociata e Varianza Spiegata), dimostrando che identifica caratteristiche con ruoli funzionali simili, non solo token di attivazione simili. Ha notevolmente avuto successo nell'abbinare sottili caratteristiche di "addizione di cifre" dove altri metodi hanno fallito.
• Compressione dei Circuiti: Il metodo ha raggiunto una maggiore accuratezza nella compressione dei circuiti di caratteristiche in super-nodi intuitivi rispetto al raggruppamento manuale e ad altre baseline. Ha identificato con successo gruppi modulari di caratteristiche (ad esempio, "testine di induzione" e "rilevamento di token") che svolgono coerentemente lo stesso ruolo funzionale attraverso circuiti diversi.
• Studio di Ablazione: Gli esperimenti hanno confermato che rappresentare le caratteristiche come distribuzioni preserva più informazioni semantiche rispetto al collassarle in un singolo vettore baricentro, anche quando proiettate sulla stessa varietà.

5. Significato e Affermazioni

Il documento afferma di fornire una metrica semantica unificata che risolve sia l'abbinamento delle caratteristiche tra livelli che la compressione dei circuiti come istanze dello stesso problema fondamentale: la stima della distanza semantica attraverso varietà.

• Scalabilità: Automatizzando il raggruppamento delle caratteristiche in super-nodi, il metodo rimuove la dipendenza dall'etichettatura umana, rendendo fattibile l'analisi di circuiti di caratteristiche su larga scala.
• Robustezza: L'approccio distribuzionale cattura il "contesto" dell'attivazione della caratteristica, permettendo un abbinamento accurato anche quando le caratteristiche evolvono significativamente tra i livelli (distribuzioni dissimili), uno scenario in cui i metodi basati su vettori falliscono.
• Interpretabilità: La capacità di comprimere automaticamente circuiti complessi in super-nodi interpretabili migliora la comprensione dei meccanismi del modello, come dimostrato negli studi di caso sui circuiti "Dallas" (geografia), "Rima" e "Opposto".

Gli autori posizionano questo lavoro come un passo fondamentale verso l'interpretabilità meccanicistica scalabile e automatizzata, superando i limiti del confronto delle caratteristiche centrato sui vettori.