Toward Identifiable Sparse Autoencoders

Technische Samenvatting: Naar Identificeerbare Sparse Autoencoders

Probleemstelling

Sparse Autoencoders (SAEs) zijn een standaardinstrument geworden voor het interpreteren en interveniëren in de interne representaties van neurale netwerken, met name Large Language Models (LLMs). Ze deconstrueren hoogdimensionele activaties naar ijle (sparse) lineaire combinaties van geleerde "concepten" (dictionary atoms). Echter, het paper identificeert een kritiek gebrek in de huidige inzet van SAEs: instabiliteit. Empirische folklore en recente theoretische werken suggereren dat verschillende trainingsruns op dezelfde data vaak resulteren in verschillende concept-dictionaries en sparse codes. Deze niet-identificeerbaarheid ondermijnt de betrouwbaarheid van SAEs voor mechanistische interpreteerbaarheid, steering en wetenschappelijke analyse, aangezien de geëxtraheerde "features" artefacten kunnen zijn van het trainingsproces in plaats van stabiele structuren in de data.

De auteurs stellen dat hoewel dictionary learning-theorie voorwaarden biedt voor identificeerbaarheid, standaard SAE-architecturen hieraan niet voldoen vanwege drie specifieke kenmerken, de zogenaamde "dark triad":

1. Bidirectionele Features: Standaard activatiefuncties (zoals ReLU of TopK) beperken coëfficiënten tot positieve waarden, waardoor de dictionary gedwongen wordt om tegengestelde concepten (bijv. $d$ en $-d$) te leren om bidirectionele signalen te representeren, wat de voorwaarden voor unieke recovery schendt.
2. Dictionary Conditioning: Standaard SAEs dwingen geen condities af (zoals de Restricted Isometry) op de geleerde dictionary, wat leidt tot scenario's waarin meerdere sparse codes dezelfde input even goed kunnen reconstrueren.
3. Encoder Expressiviteit: De bijna-lineaire encoders die gewoonlijk voor SAEs worden gebruikt, missen mogelijk de capaciteit om observaties naar de correcte sparse codes te mappen, zelfs als de dictionary goed geconditioneerd is.

Methodologie

Het paper stelt een verenigde aanpak voor om de kloof tussen de schaalbaarheid van SAEs en de theoretische garanties van dictionary learning te overbruggen. De auteurs introduceren de Identifiable Sparse Autoencoder (iSAE), die de standaard TopK SAE-architectuur en trainingsprocedure aanpast via drie specifieke interventies:

1. Bidirectionele Features (AbsTopK)

Om de beperking tot positieve coëfficiënten aan te pakken, vervangen de auteurs de standaard TopK-activatie door AbsTopK. Deze functie selecteert de $k$ coëfficiënten met de grootste absolute waarden, waardoor zowel positieve als negatieve loadings mogelijk zijn. Dit verdubbelt effectief de representatiecapaciteit van de dictionary zonder het aantal parameters te verhogen en brengt de SAE in lijn met klassieke sparse coding-formuleringen die negatieve coëfficiënten toestaan.

2. Dictionary Conditioning (aRIP Regularisatie)

De auteurs introduceren een datagestuurde relaxatie van de Restricted Isometry Property (RIP), getiteld approximate RIP (aRIP).

• Theorie: Klassieke RIP vereist dat elke combinatie van $k$ atomen normen behoudt, een conditie die combinatorisch moeilijk te verifiëren is voor overcomplete, geleerde dictionaries. De auteurs stellen voor om deze eigenschap alleen af te dwingen op de unie van de daadwerkelijk geobserveerde sparsity-patronen in de data.
• Implementatie: Ze leiden een regularisator af op basis van de eigenwaarden van de Gram-submatrix van de geobserveerde supports. Deze regularisator wordt toegevoegd aan de standaard reconstructieverlies om de geleerde dictionary aan te moedigen de aRIP-conditie te voldoen, wat ervoor zorgt dat sparse codes uniek te herstellen zijn (op permutatie en teken na).

3. Encoder Expressiviteit (Multi-Step Encoding)

Om de beperking van bijna-lineaire encoders aan te pakken, integreren de auteurs een multi-step encoding schema geïnspireerd door LISTA (Gregor & LeCun, 2010). In plaats van een enkele forward pass, verfijnt de encoder de sparse code $z$ iteratief over $T$ stappen (doorgaans $T=5$) met behulp van een geleerde stapgrootte-matrix. Dit vergroot de expressiviteit van de mapping van input $x$ naar code $z$, waardoor de geamortiseerde encoder beter de oplossing van het sparse coding optimalisatieprobleem kan benaderen.

Belangrijkste Bijdragen

Theoretische Bijdragen

• Definitie van Statistische Identificeerbaarheid: Het paper definieert SAE-identificeerbaarheid niet als het terugvinden van een "ground truth" mechanisme (wat vaak onbekend is), maar als run-to-run stabiliteit. Een SAE is identificeerbaar als onafhankelijke trainingsruns op dezelfde distributie ongeveer dezelfde dictionary en sparse codes opleveren (op triviale equivarianties zoals permutatie en teken na).
• Onmogelijkheidstheorema (Theorem 3.2): De auteurs bewijzen dat het leren van een stabiele dictionary niet voldoende is voor het leren van stabiele codes. Zelfs met een perfect stabiele dictionary, als de dictionary niet voldoet aan specifieke conditioneringseigenschappen, blijft het sparse coding probleem niet-identificeerbaar.
• Near-Identifiability Theorem (Theorem 3.6): Ze bewijzen dat als een SAE een lage reconstructiefout bereikt en de dictionary aan de aRIP-conditie voldoet (met voldoende rijkdom en diversiteit in de datadistributie), de geleerde dictionary en sparse codes near-identifiable zijn.

Empirische Bijdragen

• Nieuwe Metrieken: De auteurs stellen voor om code-identificeerbaarheid direct te meten, met behulp van Intersection-over-Union (IoU) en $\ell_2$-fout op sparse patronen, naast de standaard Dictionary Cosine Similarity (DCS). Ze demonstreren dat een hoge DCS geen hoge code-identificeerbaarheid garandeert.
• Modelvarianten: Ze presenteren twee varianten:
  • iSAE: Gebruikt AbsTopK + aRIP regularisatie.
  • iSAE-ME: Gebruikt AbsTopK + aRIP regularisatie + Multi-Step Encoding.

Resultaten

Synthetische Data

• Reconstructie: Het overstappen van TopK naar AbsTopK verminderde de reconstructiefout met 30–40%.
• Identificeerbaarheid: De iSAE-ME variant bereikte bijna perfecte identificeerbaarheid in synthetische settings. De multi-step encoder was cruciaal hier, omdat deze de geamortiseerde encoder in staat stelde om de prestaties van een "oracle" solver (Orthogonal Matching Pursuit) te evenaren die toegang heeft tot de ware dictionary.
• Conditioning: De aRIP regularisator verbeterde de conditionering van de dictionary succesvol, wat noodzakelijk was om de oracle te laten excelleren in identificeerbaarheid.

Real-World Data (LLM en Vision)

• Reconstructie: Op Pythia-160M (LLM) en DINOv2 (Vision) activaties bereikte de multi-step encoder (iSAE-ME) de laagste reconstructiefout (MSE), waarmee het baselines significant versloeg.
• Identificeerbaarheid:
  • Dictionary: AbsTopK verbeterde de dictionary-stabiliteit (DCS) vergeleken met TopK.
  • Codes: De resultaten waren gemengd. Hoewel de oracle solver een verbeterde code-identificeerbaarheid liet zien met de geconditioneerde dictionary, maakte de geamortiseerde SAE-encoder dit niet altijd volledig uit. In sommige real-world gevallen maakte de multi-step encoder de code-identificeerbaarheid zelfs slechter ondanks een betere reconstructie, wat suggereert dat optimalisatiedynamiek in real-world settings complex is en de encoder mogelijk nog niet naar de optimale oplossing convergeert.
• Downstream Taken (SAEBench): Op LLM-activaties vertoonde iSAE-ME de beste prestaties in het minimaliseren van de cross-entropy loss (reconstructiefidelity). Echter, voor "Spurious Correlation Removal" (SCR) presteerde de TopK baseline het best, waarschijnlijk omdat het op nul zetten van een TopK-feature een duidelijkere interpretatie van "afwezigheid" heeft. iSAE-ME presteerde het best op Targeted Probe Perturbation (TPP).

Betekenis en Claims

Het paper claimt dat identificeerbaarheid een "minimale vereiste" is voor SAEs die worden gebruikt in toepassingen met een hoog risico, zoals model steering en oversight. Als twee SAEs getraind op dezelfde data verschillende features produceren, komt de wetenschappelijke validiteit van die features in het gedrang.

De auteurs positioneren hun werk als een stap naar het statistisch betrouwbaar maken van SAEs. Ze beweren niet het probleem te hebben opgelost van het terugvinden van de "ware" ground-truth concepten, noch beweren ze dat hun modellen identificeerbaarheid generaliseren naar distributies buiten de trainingsdata. In plaats daarvan demonstreren ze dat:

1. Standaard TopK SAEs inherent instabiel en niet-identificeerbaar zijn.
2. Door expliciet de dictionary-conditioning (via aRIP) en encoder-expressiviteit aan te pakken, men SAEs kan construeren die statistisch identificeerbaar zijn (stabiel over verschillende runs) en een superieure reconstructie bereiken.
3. De kloof tussen synthetische en real-world prestaties benadrukt dat hoewel het theoretische kader standhoudt, optimalisatie-uitdagingen in de werkelijke neurale netwerkactivaties nog steeds een significante barrière vormen voor perfecte identificeerbaarheid.

Het paper concludeert dat deze architecturale en regularisatie-wijzigingen modellen opleveren die robuuster en theoretisch beter onderbouwd zijn, wat een sterkere fundering biedt voor de interpreteerbaarheid van grootschalige neurale netwerken.