Technische Samenvatting: Semantische Optimaal Transport voor Matching van Kenmerken in Sparse Autoencoders en Compressie van Circuits

1. Probleemdefinitie

Sparse Autoencoders (SAE's) zijn een belangrijk hulpmiddel geworden voor het interpreteren van Large Language Models (LLM's) door latente ruimtes op te splitsen in interpreteerbare kenmerken. Er blijven echter twee kritieke uitdagingen bij het schalen van SAE-analyses:

1. Matching van Kenmerken over Lagen: Het identificeren van semantisch vergelijkbare kenmerken over verschillende lagen van een LLM. Bestaande methoden vertrouwen doorgaans op het vergelijken van decodervectoren. Deze aanpak faalt bij het matchen van kenmerken over niet-aaneengesloten lagen, omdat decodervectoren zijn getraind om verschillende activatiedistributies te reconstrueren, waardoor directe vectorvergelijking onnauwkeurig wordt.
2. Compressie van Kenmerkcircuits: Kenmerkcircuits (causale grafieken die kenmerken over lagen met elkaar verbinden) bevatten vaak honderden knopen, wat handmatige inspectie onpraktisch maakt. Bestaande compressiemethoden vertrouwen op het handmatig groeperen van semantisch vergelijkbare knopen tot "supernodes", een proces dat niet schaalbaar is.

Het kernprobleem dat door de auteurs wordt geïdentificeerd, is de schatting van de semantische afstand tussen kenmerken die op verschillende activatiemanifolden liggen. Bestaande benaderingen behandelen kenmerken als enkele vectoren (bijvoorbeeld decodergewichten), wat de rijke context van waar en hoe sterk een kenmerk activeert, weggooit. Dit leidt tot informatieverlies en het onvermogen om vergelijkingen over meerdere manifolden aan te pakken.

2. Methodologie

De auteurs stellen een distributief kader voor dat kenmerkvergelijking herformuleert als een Optimaal Transport (OT)-probleem.

2.1 Representatie van Kenmerken als Distributies

In plaats van een kenmerk weer te geven door een enkele decodervector, stelt de methode elk kenmerk $i$ in laag $\ell$ voor als een activatie-gewogen empirische distributie over de verborgen staten waar het activeert.

• Gegeven een corpus van tokenposities $T$, selecteert de methode de top-$K$ tokenposities waar kenmerk $i$ het sterkst activeert.
• Elke geselecteerde token $t$ krijgt een gewicht $w_{i,t}$ toegewezen dat evenredig is met de grootte van de activatie.
• Het kenmerk wordt weergegeven als een distributie $\hat{\mu}_{i,T} = \sum w_{i,t} \delta_{x_t}$, waarbij $x_t$ de verborgen staat is bij token $t$.

2.2 Gedeelde Referentieruimte

Om kenmerken over verschillende lagen te vergelijken (die op verschillende manifolden liggen), projecteert de methode deze laagspecifieke distributies naar een gedeelde referentieruimte $\mathcal{R}$.

• Voor matching over lagen is de referentieruimte doorgaans de ruimte van verborgen staten van de doel-laag. Activaties uit de bronlaag worden geprojecteerd naar de corresponderende tokenposities in de doel-laag.
• Voor circuitcompressie is de referentieruimte de concatenatie van verborgen staten over alle lagen.

2.3 Semantische Afstand via Wasserstein-afstand

Zodra geprojecteerd in de gedeelde ruimte, wordt de semantische afstand tussen twee kenmerken berekend met de Wasserstein-afstand (Earth Mover's Distance).

• De afstand $W_c(\hat{\mu}, \hat{\nu})$ meet de minimale kosten om de waarschijnlijkheidsmassa van de activatiedistributie van het ene kenmerk naar het andere te transporteren.
• De grondkosten $c(z, z')$ meten de kosten van het matchen van twee contextrepresentaties in de gedeelde ruimte.
• Deze aanpak vangt de geometrische structuur van activatiepatronen, in plaats van alleen punt-voor-punt vectorgelijkenis.

2.4 Downstreamtaken

• Matching van Kenmerken: Voor een doelkenmerk vindt de methode het bronkenmerk met de minimale Wasserstein-afstand (naaste buur).
• Circuitcompressie: De methode behandelt circuitknopen als punten in de distributieruimte en past clustering toe (bijvoorbeeld Agglomeratieve clustering) om knopen met kleine paarwijze Wasserstein-afstanden te groeperen tot supernodes, waarbij de semantische afstand binnen clusters wordt geminimaliseerd.

3. Belangrijkste Bijdragen en Theoretische Waarborgen

Het artikel biedt drie belangrijke theoretische waarborgen voor het voorgestelde kader:

1. Invariantie voor Rescaling van Activaties: De geprojecteerde kenmerkdistributie is invariant voor globale schaling van activatiewaarden. Omdat de representatie berust op de relatieve rangschikking en de proportionele distributie van activaties (Top-K-selectie en genormaliseerde gewichten), veranderen willekeurige scalair verschuivingen in activatiegroottes de semantische afstand niet.
2. Stabiliteit onder Perturbaties: De Wasserstein-afstand tussen kenmerkdistributies is stabiel onder gladde perturbaties van de distributies. De verandering in afstand wordt begrensd door de Lipschitz-constanten van de grondkosten en de Wasserstein-afstand tussen de geperturbeerde en originele distributies. Dit waarborgt robuustheid tegen ruis in de tokenselectie of het hertrainen van SAE's.
3. Herstel op Basis van Marges: De auteurs stellen voorwaarden vast waaronder de empirische matching van de naaste buur en semantische groepering de ware (oracle) matches perfect herstellen. Als de "marge" (het gat tussen de beste match en de op één na beste match) meer dan twee keer de schattingsfout bedraagt, is de empirische beslissing gegarandeerd correct.

4. Experimentele Resultaten

De methode werd geëvalueerd op SAE's getraind op GPT-2 en Gemma-2-2B, met vergelijking tegen baselines waaronder SAE Match (L2 van decodervectoren), Feature Flow (cosinus van decodervectoren), Attribution Patching en LLM-gebaseerde selectie.

• Matching van Kenmerken: De voorgestelde methode presteerde beter dan alle baselines in LLM-evaluatie (semantische gelijkenisbeoordelingen) en Matchingsnauwkeurigheid over opeenvolgende, midden-naar-laatste en eerste-naar-laatste laagopstellingen. Cruciaal behield het concurrerende prestaties op downstream-metrics (Cross-Entropy-verlies en Verklaarde Variantie), wat aantoont dat het kenmerken identificeert met vergelijkbare functionele rollen, niet alleen vergelijkbare activatietokens. Het slaagde opvallend in het matchen van subtiele "cijferoptelling"-kenmerken waar andere methoden faalden.
• Circuitcompressie: De methode behaalde een hogere nauwkeurigheid bij het comprimeren van kenmerkcircuits naar intuïtieve supernodes in vergelijking met handmatige groepering en andere baselines. Het slaagde erin modulaire kenmerkgroepen te identificeren (bijvoorbeeld "inductiekoppen" en "tokendetectie") die consistent dezelfde functionele rol vervullen in verschillende circuits.
• Ablatiestudie: Experimenten bevestigden dat het weergeven van kenmerken als distributies meer semantische informatie behoudt dan het samenvoegen ervan tot een enkele centroidvector, zelfs wanneer geprojecteerd op dezelfde manifold.

5. Betekenis en Claims

Het artikel claimt een unificerend semantisch maatstaf te bieden die zowel matching van kenmerken over lagen als circuitcompressie oplost als gevallen van hetzelfde fundamentele probleem: het schatten van semantische afstand over manifolden.

• Schaalbaarheid: Door het groeperen van kenmerken in supernodes te automatiseren, verwijdert de methode de afhankelijkheid van menselijke labeling, waardoor de analyse van grootschalige kenmerkcircuits haalbaar wordt.
• Robuustheid: De distributieve aanpak vangt de "context" van kenmerkactivatie, wat nauwkeurige matching mogelijk maakt zelfs wanneer kenmerken aanzienlijk evolueren over lagen (ongelijksoortige distributies), een scenario waarin vectorgebaseerde methoden falen.
• Interpreteerbaarheid: Het vermogen om complexe circuits automatisch te comprimeren tot interpreteerbare supernodes verbetert het begrip van modelmechanismen, zoals aangetoond in casestudies over "Dallas" (geografie), "Rijmen" en "Tegengestelde" circuits.

De auteurs positioneren dit werk als een fundamentele stap naar schaalbare, geautomatiseerde mechanistische interpretatie, die voorbijgaat aan de beperkingen van vectorcentrische kenmerkvergelijking.