Technisch overzicht: VASAE – Het benoemen van SAE-dictaatrichtingen met vocabulaire-uitgelijnde verankering

1. Probleemstelling

Sparse Autoencoders (SAEs) worden steeds vaker gebruikt om de hoogdimensionale residual streams van Transformer-modellen te ontleden in ijle codes over geleerde dictaten, wat bijdraagt aan mechanistische interpreteerbaarheid. Echter, een belangrijke beperking van standaard SAE's is dat hun geleerde kenmerken (feature directions) niet intrinsiek verbonden zijn met de token-vocabulaire van het model.

In de standaardpraktijk is feature-benaming een post hoc proces. Onderzoekers moeten de contexten inspecteren waarin specifieke features activeren of aparte automatische interpretatieprocedures uitvoeren om semantische labels toe te wijzen. Deze discrepantie ontstaat omdat de geometrie van het geleerde dictaat en de geometrie van de token-embeddings doorgaans ongerelateerd zijn; een feature kan effectief activaties reconstrueren terwijl deze ver verwijderd ligt van een token-embeddingvector, waardoor deze zonder "intrinsieke" naam op vocabulaireniveau blijft.

De paper stelt de vraag: Kunnen we SAE-features zodanig trainen dat ze flexibel genoeg blijven voor hoogwaardige reconstructie, maar geometrisch dichtbij genoeg verankerd zijn bij token-embeddings om tijdens de training intrinsieke tokennamen te verwerven, zonder dat dit ten koste gaat van de reconstructieprestaties?

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

De auteurs stellen VASAE voor, een modificatie van de standaard SAE-trainingsdoelstelling die een mechanisme van vocabulaire-uitgelijnde verankering introduceert.

Kernarchitectuur

• Input: VASAE opereert op de post-residual stream vectoren ($h^{(\ell)}$) van een Transformer-laag $\ell$.
• Encoder: Gebruikt een standaard top-k ijle encoder (lineaire mapping + ReLU + Top-K selectie) om een ijle code $z$ te produceren.
• Decoder: Gebruikt een leerbare dictaatmatrix $W_D$ (waarbij kolommen feature-vectoren $d_i$ zijn) om de residual stream te reconstrueren.
• Ankers: De methode maakt gebruik van de vaste input-token embeddings ($W_E$) van het model als referentierichtingen. Deze embeddings delen dezelfde dimensionaliteit als de residual stream (of worden geprojecteerd naar die dimensie).

De Verankeringsdoelstelling

De belangrijkste innovatie is de toevoeging van een anchor loss term aan de standaard reconstructieverliesfunctie.

1. Naderste-Token Uitlijning: Voor elke geleerde feature $d_i$ berekent de methode de cosinusovereenkomst met alle vaste token-embeddings $w_v$.
   $$s_i = \max_{v \in V} \cos(d_i, w_v)$$
   De token $v^*_i$ die deze maximumwaarde bereikt, definieert de intrinsieke tokennaam voor feature $i$.
2. Anchor Loss: De trainingsdoelstelling moedigt features aan om dicht bij hun dichtstbijzijnde token-embedding te blijven:
   $$L_{anchor} = -\frac{1}{d_{sparse}} \sum_{i=1}^{d_{sparse}} s_i$$
   Het minimaliseren van deze loss maximaliseert de cosinusovereenkomst tussen features en hun dichtstbijzijnde vocabulaire-tokens.
3. Totale Loss: De gecombineerde doelstelling is:
   $$L_{vasae} = L_{recon} + \lambda_{anchor} L_{anchor}$$
   waarbij $\lambda_{anchor}$ een hyperparameter is die de sterkte van de uitlijningsrestrictie controleert.

Onderscheid van Hard-Tying

De auteurs onderscheiden VASAE expliciet van een "hard-tied" decoder baseline, waarbij de decodergewichten exact gelijk zijn aan de token-embeddingmatrix ($W_D = W_E$). Hard-tying dwingt features om exacte token-embeddings te zijn, wat volgens de paper onvoldoende is voor hoogwaardige reconstructie. VASAE houdt het dictaat leerbaar, maar trekt het richting de geometrie van de vocabulaire via de anchor loss.

3. Experimentele Opstelling

De methode werd geëvalueerd op twee decoder-only Transformer-modellen:

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

Baselines:

1. Standaard SAE: Getraind met enkel reconstructieverlies.
2. Hard-Tied Decoder: Decodergewichten vastgezet op token-embeddings; alleen de encoder wordt getraind.

Metrieken:

• Reconstructie: Gemiddelde kwadratische fout (MSE), verklaarde variantie (VE) en cross-entropie (CE) herstel na substitutie van de residual stream.
• Uitlijning: Geometrische feature-uitlijningsgraad (percentage features met een uitlijningsscore $s_i \ge 0.8$) en vocabulaire-dekking.
• Casestudies: Kwalitatieve inspectie van intrinsieke tokennamen toegewezen aan actieve features in specifieische contexten.

4. Belangrijkste Resultaten

Reconstructiekwaliteit

VASAE bereikt reconstructiemetrieken die vergelijkbaar zijn met de Standaard SAE, en presteert significant beter dan de Hard-Tied Decoder baseline.

• GPT-2-small: Zowel VASAE als de Standaard SAE bereiken een verklaarde variantie (VE) van 0.965 en een CE-herstel van 0.975.
• Llama-3.1-8B: Zowel VASAE als de Standaard SAE bereiken een VE van 0.931 en een CE-herstel van 0.906.
• Hard-Tied Baseline: Presteert slecht (bijv. negatieve VE in GPT-2-small), wat bevestigt dat het afdwingen van exacte token-embeddings de reconstructiecapaciteit vernietigt.

Geometrische Uitlijning

• GPT-2-small: Met $\lambda_{anchor} = 10^{-4}$ overschrijden 89–94% van de features in lagen 0–10 de uitlijningsdrempel ($s_i \ge 0.8$). De uitlijning is sterk en consistent in ondiepe en middelste lagen.
• Llama-3.1-8B: De uitlijning is diepte-afhankelijk. Bij een hogere coëfficiënt ($\lambda_{anchor} = 5 \times 10^{-3}$) bereikt de ondiepe laag (L0) een uitlijning van 92.8%. Echter, de uitlijning neemt aanzienlijk af in diepere lagen (bijv. L31), wat suggereert dat vocabulaire-verankering compatibeler is met residual states die dichter bij de embedding-ruimte liggen.

Casestudies

Kwalitatieve analyse van intrinsieke tokennamen onthult semantische coherentie:

• In GPT-2-small worden features die actief zijn in een zin over "Baker Street" toegewezen namen zoals located, Street, corner, en around.
• Aanvullende voorbeelden tonen features die uitlijnen met bijvoeglijke naamwoorden/bijwoorden (bijv. perfect, feasible), termen gerelateerd aan prijzen (bijv. winner, praise), en termen voor zelfintroductie (bijv. fan, name).
• In Llama-3.1-8B vertoont de ondiepe laag een vergelijkbare uitlijning met locaties, terwijl diepere lagen minder consistent zijn, wat de kwantitatieve uitlijningsmetrieken weerspiegelt.

5. Bijdragen en Betekenis

De paper levert drie primaire bijdragen:

1. Introductie van VASAE: Een zachte vocabulaire-verankeringsdoelstelling die geleerde SAE-decoderrichtingen tijdens de training de dichtstbijzijnde tokennamen toeweist, in plaats van post-hoc.
2. Evaluatie van Verankering: Het aantonen dat vocabulaire-uitgelijnde verankering een reconstructiekwaliteit behoudt die vergelijkbaar is met standaard SAE's, terwijl het bewijst dat het hard-tyen van de decoder aan token-embeddings onvoldoende is voor deze taak.
3. Geometrische Analyse: Het analyseren van de uitlijning van embeddings in de feature-ruimte en de vocabulaire-dekking, wat bewijs levert dat intrinsieke tokennamen fungeren als geometrie-afgeleide identificatoren.

Betekenis:
De auteurs stellen dat VASAE bestaande post-hoc interpretatiemethoden aanvult door geleerde features te verbinden met intrinsieke tokennamen tijdens de training. Dit biedt een "vocabulaire-geïndexeerde handvat" voor features, wat een startpunt biedt voor inspectie.

Cruciaal is dat de auteurs een bescheiden standpunt innemen over de implicaties van deze namen:

• Geen Verklaringen: Een intrinsieke tokennaam is geen volledige verklaring van een feature. Het legt niet uit waarom een feature activeert of wat de causale rol ervan is in de berekening van het model.
• Indexeringsinstrument: De namen dienen als voorlopige indexeringsinstrumenten. Echt begrip vereist verdere analyse van hoe genoemde features deelnemen aan circuits, hoe ze verbinding maken met andere componenten, en hoe interventies op hen het gedrag beïenvloeden.
• Beperkingen: De uitlijning is het meest robuust in GPT-2-small en de ondiepe/middelste lagen van Llama-3.1-8B. De methode verankert momenteel alleen aan input-embeddings; voor modellen met niet-gekoppelde input/output-embeddings blijft de keuze van de anker-matrix een openstaande vraag.

Samenvattend slaagt VASAE erin de kloof te overbruggen tussen de geometrische ruimte van geleerde features en de vocabulaire-ruimte van het model, waardoor features intrinsieke tokennamen kunnen dragen zonder hun nut voor reconstructie te verliezen.