Technische Samenvatting: Rational Sparse Autoencoder (RSAE)

Probleemstelling

Sparse Autoencoders (SAEs) zijn het standaardinstrument voor mechanistische interpreteerbaarheid in large language models (LLMs), waarbij interne activaties worden gedecomposeerd in ijle (sparse) lineaire combinaties van overcomplete feature-dictionaries. Echter, bestaande SAE-families vertrouwen op vaste, niet-gladde encoder-nietlineariteiten—specifiek ReLU (met $\ell_1$-regularisatie), TopK en JumpReLU. Deze vaste activatieprimitieven introduceren specifieke pathologieën:

• ReLU SAEs: Lijden aan magnitude-krimp van actieve features en een hardnekkige populatie van "dode" latents.
• TopK SAEs: Vervangen zachte straffen door harde cardinaliteitsrestricties, wat de gradiëntstroom door inactieve features doorbreekt en hulpbronnen voor revival vereist.
• JumpReLU SAEs: Introduceren leerbare drempelwaarden, maar vereisen continue-relaxatie surrogaten voor back-propagation door hun indicator-gates.

Deze hardgecodeerde mechanismen beperken het vermogen van het model om zich aan te passen aan de geobserveerde pre-activatiegeometrie, wat potentieel de afweging tussen reconstructietrouw en ijleheid (sparsity) verstoort.

Methodologie

De auteurs stellen de Rational Sparse Autoencoder (RSAE) voor, die de vaste encoder-activatie vervangt door een leerbare rationale functie. De RSAE behoudt de standaard SAE-structuur (een affine pre-activatielaag gevolgd door een ijle activatieblok en een lineaire decoder), maar wijzigt de activatiefunctie $\phi(\cdot)$ naar een leerbare rationale functie die elementgewijs wordt toegepast.

Architectuur

De RSAE-activatie wordt gedefinieerd als:
$$\phi(h) = C_{out} \cdot r_{(a,b)}\left(\frac{h}{C_{in}}\right)$$
waarbij $h$ de pre-activatievector is, en $C_{in}$ en $C_{out}$ leerbare schalingparameters zijn, en $r_{(a,b)}(t) = P(t)/Q(t)$ een rationale functie is met leerbare coëfficiënten $(a, b)$. De schalingparameters mappen de pre-activatieverdeling naar het ontwerpinterval $[-1, 1]$ van de rationale functie en mappen de output terug naar de verwachte feature-magnitude.

Theoretische Fundering

De paper vestigt een expressieve asymmetrie tussen rationale functies en de stuksgewijze affine gates die in huidige SAEs worden gebruikt:

1. Bovengrens (Rationale Approximatie): Met behulp van Zolotarev sign-functies bewijzen de auteurs dat leerbare rationale activaties ReLU, JumpReLU en geleverde thresholded TopK-gates kunnen benaderen met een polylogarithmische grootte in de doelnauwkeurigheid $\epsilon$ (specifiek $O(\log(1/\epsilon)\log(1/\delta))$ voor discontinue gates met marge $\delta$).
2. Ondergrens (Piecewise-Affine Limiet): Omgekeerd demonstreren zij dat het benaderen van bepaalde $O(1)$-parameter rationale target-maps met scalaire output single-layer stuksgewijze affine encoders (ReLU/JumpReLU/TopK) $\Omega(\epsilon^{-1/2})$ geactiveerde coördinaten vereist. Dit suggere houdt aan dat rationale activaties een efficiëntere representatieklasse bieden voor de SAE-encoder.

Trainingspipeline

De RSAE wordt geconstrueerd via een tweestaps-pipeline die ontworpen is om vooraf getrainde "teacher" SAEs te upgraden:

1. Initialisatie:
   • Remez Exchange: Een rationale functie wordt gefit aan de activatie-primitive van de teacher (ReLU, JumpReLU of TopK) op een synthetische dichte grid met behulp van het gerelaxeerde Remez exchange algoritme om de $L_\infty$ fout te minimaliseren.
   • Kalibratie: De rationale coëfficiënten en schalingparameters ($C_{in}, C_{out}$) worden aangepast aan de werkelijke pre-activatieverdeling van de teacher door de kwadratische fout tussen de rationale output en de teacher output te minimaliseren.
   • Gewicht-overerving: De encoder- en decodergewichten worden letterlijk gekopieerd van de vooraf getrainde teacher.
2. Fine-tuning: Alle parameters (gewichten, biases en rationale coëfficiënten) worden ontdoordrukt (unfrozen) en gezamenlijk geoptimaliseerd onder de standaard $\ell_1$-geregulariseerde reconstructie-objectief.

Belangrijkste Bijdragen

1. RSAE Architectuur: Introductie van een sparse autoencoder met een leerbare rationale activatie, geworteld in benaderingstheorie die is afgestemd op de ondiepe SAE-encoder.
2. Theoretische Expressiviteitskloof: Bewijs dat rationale activaties compacte benaderingen bieden voor vaste SAE-gates, terwijl stuksgewijze affine encoders aanzienlijk meer geactiveerde coördinaten nodig hebben om eenvoudige rationale targets te benaderen.
3. Tweestaps Trainingsalgoritme: Een praktische initialisatie- en fine-tuning procedure die het mogelijk maakt om RSAEs als "drop-in" te gebruiken om vooraf getrainde SAEs te vervangen met verwaarloosbare overhead.
4. Empirische Validatie: Uitgebreide evaluatie over drie open-weight taalmodellen (GPT-2 small, Pythia-160m, Gemma-2-2B) en alle drie de baseline activatie-families.

Resultaten

De auteurs evalueren de RSAE tegenover getrainde baselines op residual-stream activaties. De resultaten wijzen op strikte verbeteringen na fine-tuning:

• Reconstructie-metrieken: RSAE bereikt consistent een lagere Mean Squared Error (MSE), een lagere $\ell_0$ (ijleheid) en een hoger aandeel "levende" features vergeleken met de teacher-modellen over alle geteste configuraties heen.
• Downstream Metrieken: RSAE verbetert het downstream gedrag, waarbij een verminderde cross-entropy degradatie ($\Delta CE$) en een hogere loss recovery ($LR$) wordt getoond bij het intercepten van de residual stream.
• Interpreteerbaarheid: Sparse probing experimenten op Pythia-160m tonen aan dat RSAE de feature-niveau interpreteerbaarheid behoudt. De full-dictionary probe nauwkeurigheid op acht binaire classificatietaken blijft vergelijkbaar met de teacher, zonder significante degradatie.
• Efficiëntie: De upgrade voegt slechts een handvol scalaire parameters toe per autoencoder. De gehele pipeline (initialisatie + 2.000 fine-tuning stappen) draait binnen minuten op een enkele consumenten-GPU (NVIDIA RTX 5090).
• Robuustheid: Winsten zijn consistent over verschillende host-modellen, baseline activatie-families en een breed bereik van baseline ijleheid-niveaus ($\ell_0 \in [51, 683]$).

Betekenis en Claims

De paper claimt dat de RSAE een "drop-in" modificatie biedt voor vooraf getrainde SAEs die de modelgetrouwheid strikt verbetert bij een gelijkmatige ijleheid zonder de interpreteerbaarheid op te offeren. Door vaste nietlineariteiten te vervangen door leerbare rationale functies, past de RSAE zich aan de specifieke geometrie van de pre-activaties in het host-model aan, waarmee de beperkingen van hardgecodeerde activatie-primitieven worden overwonnen.

De auteurs benadrukken dat deze verbeteringen worden bereikt met verwaarloosbare computationele kosten (minuten op een enkele GPU) en minimale parameter-overhead. Zij positioneren RSAE als een universele upgrade-strategie voor bestaande SAEs, in staat om zowel de reconstructiekwaliteit als de downstream taakprestaties te verbeteren terwijl de monosemantische features vereist voor mechanistische interpreteerbaarheid behouden blijven. De paper merkt op dat hoewel de evaluatie drie modellen en drie activatie-families dekt, de aanpak algemeen is en potentieel uitgebreid kan worden naar andere frontier-modellen of architecturale varianten (bijv. Gated, BatchTopK, Matryoshka SAEs).