Riepilogo Tecnico: VASAE – Naming SAE Dictionary Directions with Vocabulary-Aligned Anchoring

1. Problema

Gli Autoencoder Sparsi (SAE) sono sempre più utilizzati per decomporre gli stream residui ad alta dimensionalità dei modelli Transformer in codici sparsi su dizionari appresi, facilitando l'interpretabilità meccanicistica. Tuttavia, un limite significativo dei comuni SAE è che le loro caratteristiche apprese (direzioni del dizionario) non sono intrinsecamente connesse al vocabolario di token del modello.

Nella pratica standard, il naming delle feature è un processo post hoc. I ricercatori devono ispezionare i contesti in cui specifiche feature si attivano o eseguire procedure di interpretazione automatizzate separate per assegnare etichette semantiche. Questa disconnessione deriva dal fatto che la geometria del dizionario appreso e la geometria degli embedding dei token sono tipicamente non correlate; una feature può ricostruire efficacemente le attivazioni pur trovandosi lontana da qualsiasi vettore di embedding di un token, risultando priva di un nome "intrinseco" a livello di vocabolario.

La ricerca pone questa domanda: Possiamo addestrare le feature SAE in modo che rimangano abbastanza flessibili per una ricostruzione di alta qualità, ma siano geometricamente ancorate abbastanza vicino agli embedding dei token da acquisire nomi di token intrinseci durante l'addestramento, senza sacrificare le prestazioni di ricostruzione?

2. Metodologia: Vocabulary-Aligned Sparse Autoencoder (VASAE)

Gli autori propongono VASAE, una modifica all'obiettivo di addestramento standard degli SAE che introduce un meccanismo di ancoraggio allineato al vocabolario.

Architettura Core

• Input: VASAE opera sugli vettori post-stream residuo ($h^{(\ell)}$) di un livello Transformer $\ell$.
• Encoder: Utilizza un encoder sparso top-k standard (mappa lineare + ReLU + selezione Top-K) per produrre un codice sparso $z$.
• Decoder: Utilizza una matrice di dizionario apprendibile $W_D$ (dove le colonne sono i vettori delle feature $d_i$) per ricostruire lo stream residuo.
• Ancore (Anchors): Il metodo utilizza gli embedding fissi dei token di input del modello ($W_E$) come direzioni di riferimento. Questi embedding condividono la stessa dimensionalità dello stream residuo (o vengono proiettati per farlo).

L'Obiettivo di Ancoraggio

L'innovazione chiave è l'aggiunta di un termine di perdita di ancoraggio all'obiettivo standard di ricostruzione.

1. Allineamento al Token Più Vicino: Per ogni feature appresa $d_i$, il metodo calcola la similarità del coseno con tutti gli embedding dei token fissi $w_v$.
   $$s_i = \max_{v \in V} \cos(d_i, w_v)$$
   Il token $v^*_i$ che realizza questo massimo definisce il nome del token intrinseco per la feature $i$.
2. Perdita di Ancoraggio (Anchor Loss): L'obiettivo di addestramento incoraggia le feature a rimanere vicine al loro embedding di token più prossimo:
   $$L_{anchor} = -\frac{1}{d_{sparse}} \sum_{i=1}^{d_{sparse}} s_i$$
   Minimizzare questa perdita massimizza la similarità del coseno tra le feature e i loro token di vocabolario più vicini.
3. Perdita Totale: L'obiettivo combinato è:
   $$L_{vasae} = L_{recon} + \lambda_{anchor} L_{anchor}$$
   dove $\lambda_{anchor}$ è un iperparametro che controlla la forza del vincolo di allineamento.

Distinzione dal Hard-Tying

Gli autori distinguono esplicitamente VASAE da una baseline "hard-tied" del decoder, dove i pesi del decoder sono fissati per essere esattamente uguali alla matrice di embedding dei token ($W_D = W_E$). Il hard-tying forza le feature a essere esatti embedding di token, il che, come ipotizzato nel paper, è insufficiente per una ricostruzione di alta qualità. VASAE mantiene il dizionario apprendibile, ma lo tira verso la geometria del vocabolario tramite la perdita di ancoraggio.

3. Setup Sperimentale

Il metodo è stato valutato su due modelli Transformer decoder-only:

• GPT-2-small: 12 layer, $d_{model}=768$, dimensione del vocabolario $|V|=50.257$.
• Llama-3.1-8B: 32 layer, $d_{model}=4096$, dimensione del vocabolario $|V|=128.256$.

Baseline:

1. Standard SAE: Addestrato solo con perdita di ricostruzione.
2. Hard-Tied Decoder: Pesi del decoder fissati agli embedding dei token; viene addestrato solo l'encoder.

Metriche:

• Ricostruzione: Errore Quadratico Medio (MSE), Varianza Spiegata (VE) e recupero della Cross-Entropy (CE) dopo la sostituzione dello stream residuo.
• Allineamento: Tasso di allineamento geometrico delle feature (percentuale di feature con punteggio di allineamento $s_i \ge 0.8$) e copertura del vocabolario.
• Case Studies: Ispezione qualitativa dei nomi dei token intrinseci assegnati alle feature attive in specifici contesti.

4. Risultati Chiave

Qualità della Ricostruzione

VASAE raggiunge metriche di ricostruzione comparabili allo Standard SAE, superando significativamente la baseline Hard-Tied Decoder.

• GPT-2-small: Sia VASAE che lo Standard SAE raggiungono una Varianza Spiegata (VE) di 0.965 e un recupero della CE di 0.975.
• Llama-3.1-8B: Sia VASAE che lo Standard SAE raggiungono una VE di 0.931 e un recupero della CE di 0.906.
• Baseline Hard-Tied: Si comporta scarsamente (ad esempio, VE negativa in GPT-2-small), confermando che forzare le feature a essere esatti embedding di token distrugge la capacità di ricostruzione.

Allineamento Geometrico

• GPT-2-small: Con $\lambda_{anchor} = 10^{-4}$, l'89–94% delle feature nei layer 0–10 supera la soglia di allineamento ($s_i \ge 0.8$). L'allineamento è forte e costante nei layer superficiali e medi.
• Llama-3.1-8B: L'allineamento dipende dalla profondità. Con un coefficiente più alto ($\lambda_{anchor} = 5 \times 10^{-3}$), il layer superficiale (L0) raggiunge il 92.8% di allineamento. Tuttavia, l'allineamento diminuisce significativamente nei layer più profondi (es. L31), suggerendo che l'ancoraggio al vocabolario è più compatibile con gli stati residui vicini allo spazio degli embedding.

Case Studies

L'analisi qualitativa dei nomi dei token intrinseci rivela coerenza semantica:

• In GPT-2-small, le feature attive in una frase riguardante "Baker Street" ricevono nomi intrinseci come located, Street, corner e around.
• Esempi aggiuntivi mostrano feature allineate ad aggettivi/avverbi (es. perfect, feasible), termini legati ai premi (es. winner, praise) e termini di auto-presentazione (es. fan, name).
• In Llama-3.1-8B, lo strato superficiale mostra un simile allineamento legato alla posizione, mentre i layer profondi mostrano meno consistenza, riflettendo le metriche di allineamento quantitative.

5. Contributi e Significato

Il paper apporta tre contributi primari:

1. Introduzione di VASAE: Un obiettivo di ancoraggio morbido al vocabolario che assegna ai direzioni del decoder SAE apprese i nomi dei token più vicini durante l'addestramento, anziché post-hoc.
2. Valutazione dell'Ancoraggio: Dimostrare che l'ancoraggio allineato al vocabolario preserva una qualità di ricostruzione comparabile agli SAE standard, provando al contempo che il hard-tying del decoder agli embedding dei token è insufficiente per questo compito.
3. Analisi Geometrica: Analizzare l'allineamento delle feature nello spazio degli embedding e la copertura del vocabolario, fornendo prove che i nomi dei token intrinseci agiscono come identificatori derivati dalla geometria.

Significato:
Gli autori sostengono che VASAE integri gli esistenti metodi di interpretazione post-hoc collegando le feature apprese ai nomi dei token intrinseci durante l'addestramento. Questo fornisce un "handle" (maniglia) indicizzato al vocabolario per le feature, offrendo un punto di partenza per l'ispezione.

Fondamentalmente, il paper mantiene una posizione modesta sulle implicazioni di questi nomi:

• Non sono Spiegazioni: Un nome di token intrinseco non è una spiegazione completa della feature. Non spiega perché una feature si attiva o il suo ruolo causale nel calcolo del modello.
• Strumento di Indicizzazione: I nomi servono come strumenti preliminari di indicizzazione. La vera comprensione richiede un'ulteriore analisi di come le feature nominate partecipano ai circuiti, come si connettono ad altri componenti e come le interazioni su di esse influenzano il comportamento.
• Limitazioni: L'allineamento è più robusto in GPT-2-small e nei layer superficiali/medi di Llama-3.1-8B. Il metodo attualmente ancora solo agli embedding di input; per i modelli con embedding input/output non condivisi (untied), la scelta della matrice di ancoraggio rimane una questione aperta.

In sintesi, VASAE colma con successo il divario tra lo spazio geometrico delle feature apprese e lo spazio del vocabolario del modello, consentendo alle feature di portare nomi di token intrinseci senza compromettere la loro utilità per la ricostruzione.