Technische Samenvatting: Geresidualiseerde Sparse Auto-encoders voor Interventies in Multi-Layer Transformers

Probleemstelling

Sparse Auto-encoders (SAE's) zijn een primair instrument voor mechanistische interpreteerbaarheid, waarbij ze verborgen toestanden van transformers ontleden in schaarse, interpreteerbare kenmerken. De standaardpraktijk traint SAE's echter onafhankelijk op de ruwe activaties van elke laag. Deze aanpak negeert een fundamenteel structureel kenmerk van transformers: de residual stream is sterk gekoppeld over de diepte heen. Een aanzienlijk deel van de activatie van een bepaalde laag bestaat uit informatie die uit eerdere lagen wordt doorgegeven of lichtjes wordt getransformeerd.

Het trainen van onafhankelijke SAE's op ruwe activaties creëert twee praktische problemen voor interventies op meerdere lagen:

1. Capaciteitsredundantie: Verschillende woordenboeken per laag verspillen capaciteit door dezelfde voorspelbare, doorgegeven informatie opnieuw te leren.
2. Onvoorspelbare Interacties: Wanneer meerdere lagen gelijktijdig worden vervangen (bijvoorbeeld via kenmerkablatie of reconstructie), zijn de interacties tussen de lagen niet additief. Als dezelfde informatie in meerdere woordenboeken wordt vertegenwoordigd, kan gelijktijdige vervanging ertoe leiden dat kenmerken dupliceren, elkaar opheffen of activaties buiten de trainingsverdeling duwen, wat leidt tot cumulatieve fouten en een verslechterde modelprestatie (cross-entropy).

Methodologie: Geresidualiseerde Sparse Auto-encoders (ReSAE's)

De auteurs stellen Geresidualiseerde Sparse Auto-encoders (ReSAE's) voor om laag-specifieke informatie te ontkoppelen van voorspelbare cross-layer structuren. De methode omvat drie belangrijke fasen:

1. Paarsgewijze Affiene Residualisatie

Voor een geselecteerde subset van lagen $L = (\ell_1, \ell_2, \dots, \ell_M)$ past de methode een affiene afbeelding toe tussen opeenvolgende geselecteerde lagen. Voor een paar $(\ell_m, \ell_{m+1})$ wordt een affiene transformatie $A_m h^{(\ell_m)} + c_m$ geleerd via ridge-geregulariseerde gewone kleinste-kwadratenmethode om de activatie van de latere laag te voorspellen op basis van de eerdere laag.
De geresidualiseerde activatie $r^{(\ell_{m+1})}$ wordt vervolgens gedefinieerd als het onderdeel van de activatie van de latere laag dat niet wordt verklaard door deze affiene afbeelding:
$$r^{(\ell_{m+1})} = h^{(\ell_{m+1})} - (A_m h^{(\ell_m)} + c_m)$$
De eerste ankerlaag wordt gecentreerd rond het empirische gemiddelde.

2. Onafhankelijke Per-Lag SAE-training

Elke SAE wordt getraind op het residu in plaats van de ruwe activatie.

• Invoer: De residual activatie wordt genormaliseerd (via per-blok RMS-normalisatie) en ingevoerd in een standaard BatchTopK SAE.
• Doel: De SAE minimaliseert de reconstructieverlies van het residudoelwit.
• Onafhankelijkheid: Gradienten stromen niet tussen SAE's van verschillende blokken; elke SAE leert alleen de "nieuwe" informatie die op zijn specifieke laag wordt geïntroduceerd.

3. Reconstructie in Originele Ruimte en Online Interventie

Om ReSAE's te evalueren met standaardprotocollen, wordt het gereconstrueerde residu via de gefitte affiene keten teruggekaatst naar de ruimte van de originele activatie.

• Reconstructie: $\hat{h}^{(\ell_{m+1})} = A_m \hat{h}^{(\ell_m)} + c_m + S^{-1}_{m+1} \hat{z}^{(m+1)}$, waarbij $\hat{z}$ de SAE-reconstructie van het residu is.
• Online Interventie: Tijdens een forward pass met hooks wordt de residual invoer voor laag $\ell_m$ berekend door de voorspelde bijdrage van de eerder gereconstrueerde laag $\hat{h}^{(\ell_{m-1})}$ af te trekken van de huidige binnenkomende activatie. Dit zorgt ervoor dat de SAE op laag $\ell_m$ een invoerverdeling ziet die vergelijkbaar is met zijn trainingsdoelwit, zelfs als upstream lagen zijn gewijzigd.

Belangrijkste Bijdragen

1. ReSAE-architectuur: Introductie van een trainingsparadigma waarbij SAE's worden getraind op het onderdeel van de activatie dat niet lineair voorspelbaar is vanuit een vorige geselecteerde laag, met reconstructies samengesteld via een affiene keten.
2. Verbeterde Multi-Layer Interventie: Aantonen dat ReSAE's, ondanks het reconstrueren van minder ruwe variantie, de cross-entropy van de transformer trouwer herstellen onder gelijktijdige vervanging van meerdere lagen in vergelijking met ruwe SAE's.
3. Verminderde Redundantie: Bewijs dat residualisatie meer orthogonale decoder-richtingen produceert, wat aangeeft dat capaciteit wordt gericht op laag-lokale informatie in plaats van gedupliceerde cross-layer structuren.
4. Uitgebreide Evaluatie: Uitgebreide benchmarks op Pythia-1.4B en Gemma-2-9B die reconstructiekwaliteit, cross-entropy-verslechtering, decoder-ähnelijkheid, schaarse probing, gerichte probe-perturbatie (TPP) en verwijdering van schijnbare correlaties (SCR) bestrijken.

Resultaten

Experimenten zijn uitgevoerd op Pythia-1.4B (lagen 0, 6, 12, 18) en Gemma-2-9B (verschillende laagafstanden).

• Reconstructie versus Functionaliteit: ReSAE's vertonen een lagere Verklarende Variantie (EV) in de ruimte van de originele activatie in vergelijking met ruwe SAE's, omdat ze opzettelijk lineair voorspelbare structuren verwerpen. Deze lagere EV correleert echter met een lagere cosinus-ähnelijkheid van de decoder (minder redundantie).
• Herstel van Cross-Entropy: Onder sequentiële online vervanging van meerdere lagen (waarbij fouten zich voortplanten) presteren ReSAE's aanzienlijk beter dan ruwe SAE's in het herstellen van de cross-entropy van het model, vooral bij hogere sparsiteitsniveaus ($k \ge 160$). In teacher-forced settings (geen foutvoortplanting) bereiken ReSAE's consistent de beste cross-entropy-herstel.
• Overinteractie: Ruwe SAE's vertonen grote negatieve "overinteractie" (niet-additieve verslechtering) wanneer meerdere lagen worden vervangen. ReSAE's verminderen deze metriek naar nul, wat aangeeft dat hun vervangingen per laag voorspelbaarder samenkomen.
• Kwaliteit van Kenmerken:
  • Schaarse Probing & TPP: ReSAE's presteren over het algemeen beter dan ruwe SAE's in de nauwkeurigheid van schaarse probing en gerichte probe-perturbatie, wat suggereert dat de geleerde kenmerken nuttiger zijn voor lineaire uitlezing en selectieve ablatie.
  • Verwijdering van Schijnbare Correlaties (SCR): ReSAE's presteren hier niet uniform beter dan ruwe SAE's; in sommige settings (met name per-laag gezamenlijke ablaties op Pythia) scoren ruwe SAE's hoger. Dit suggereert dat geresidualiseerde kenmerken minder selectief kunnen zijn voor het verwijderen van specifieke schijnbare attributen zonder taaksignalen te beschadigen.
• Schaal en Afstand: De voordelen van ReSAE's blijven bestaan in het grotere Gemma-2-9B-model over verschillende laagafstanden, met consistente verminderingen in decoder-redundantie en verbeteringen in cross-entropy-herstel.

Betekenis en Claims

Het artikel betoogt dat het behandelen van de transformer residual stream als een gekoppeld object in plaats van een verzameling geïsoleerde lagen cruciaal is voor effectieve interventies op meerdere lagen.

• Capaciteitsallocatie: Door lineair voorspelbare cross-layer structuren te verwijderen, dwingen ReSAE's het schaarse woordenboek om zich te richten op de onderdelen van de activatie die het meest relevant zijn voor de downstream-berekening van het model op die specifieke laag.
• Stabiliteit van Interventies: De methode biedt een stabielere standaard voor interventies op meerdere lagen, waardoor de cumulatieve fouten en onvoorspelbare interacties worden verminderd die standaard SAE-implementaties teisteren.
• Afweging: De auteurs erkennen dat ReSAE's geen universele verbetering zijn voor alle interpreteerbaarheidstaken (bijvoorbeeld de SCR-resultaten zijn gemengd). Echter, voor het specifieke doel van vervanging op meerdere lagen en interventie op basis van kenmerken, is de afweging van lagere ruwe variantie voor hogere functionele trouw en verminderde redundantie gunstig.

Het werk suggereert dat het verwijderen van voorspelbare structuren via affiene regressie een nuttige voorverwerkingsstap is voor SAE's in multi-layer settings, hoewel de optimale keuze van predictor (affien versus niet-lineair) en laagselectie een open ontwerpruimte blijft.