Technische Samenvatting: Control Reinforcement Learning (CRL)

1. Probleemstelling

Sparse Autoencoders (SAE's) hebben zich ontwikkeld tot een krachtig hulpmiddel voor het ontleden van activaties van taalmodellen in interpreteerbare, voor mensen leesbare kenmerken. Echter, bestaande analysemethoden zijn grotendeels statisch: ze identificeren welke kenmerken op specifieke inputs activeren, maar geven niet aan welke kenmerken, wanneer versterkt, daadwerkelijk de modeloutput veranderen. Huidige technieken voor activatie-engineering vertrouwen vaak op contrastieve datasets voor kenmerkselectie of uitgebreide opslag van activaties voor coëfficiëntoptimalisatie, en missen adaptieve feedbackmechanismen om stuurstrategieën aan te passen op basis van de kwaliteit van de generatie.

Het artikel adresseert de kloof tussen statische kenmerkidentificatie en dynamische, interpreteerbare interventie. De kernuitdaging is het ontwikkelen van een raamwerk dat niet alleen SAE-kenmerken selecteert om een model te sturen, maar ook token-niveau interventielogboeken levert die uitleggen waarom een kenmerk werd geselecteerd en hoe het de output beïnvloedde, waardoor sturen wordt omgezet in een instrument voor mechanische interpretatie.

2. Methodologie: Control Reinforcement Learning (CRL)

De auteurs stellen Control Reinforcement Learning (CRL) voor, een raamwerk dat de controle van transformerrepresentaties behandelt als een Markov Decision Process (MDP) over SAE-kenmerken.

2.1. MDP Formulering

• State ($S$): De observatie is de activatie van de residustroom ($x \in \mathbb{R}^d$) op een specifieke laag $\ell$ voor de huidige tokenpositie.
• Action ($A$): Een binaire vector die één of meer SAE-kenmerken selecteert om te versterken. Het beleid selecteert de top-$k$ kenmerken (typisch $k=1$ voor interpreteerbaarheid).
• Transition ($P$): Deterministisch, bepaald door de forward pass van de transformer met de toegepaste interventie van het geselecteerde kenmerk.
• Reward ($R$): Taakspecifieke beloningen gebaseerd op de kwaliteit van de uiteindelijke output (bijvoorbeeld correctheid in MMLU, weigering in veiligheidstaken).

2.2. CRL-Token en Kenmerksturing

Bij elke generatiestap $t$ voert het raamwerk interventies op token-niveau uit:
$$\tilde{x}_t = x_t + c \cdot a_t W_{dec}$$
Waarbij $a_t$ de sparse binaire selectievector is, $c$ de stuurcoëfficiënt is, en $W_{dec}$ de SAE-decoder is. Het beleidsnetwerk $\pi_\theta$ mapt activaties van de residustroom naar logits voor kenmerkselectie, geïmplementeerd als een 2-laags MLP.

2.3. Adaptieve Kenmerkmaskering (AFM)

Om te voorkomen dat het beleid instort tot een smalle set van frequent actieve kenmerken, introduceert CRL Adaptive Feature Masking.

• Een per-steksel booleaan masker beperkt de keuzes voor kenmerkselectie.
• Het masker begint met een kleine subset van kenmerken en breidt progressief uit naarmate de generatie vordert, waarbij kenmerken worden opgenomen die in eerdere stappen van nature werden geactiveerd.
• Dit mechanisme stimuleert diverse kenmerkontdekking terwijl de interpreteerbaarheid van interventies met één kenmerk behouden blijft.

2.4. Optimalisatie

Het beleid en een criticusnetwerk (dat stateschatten schat) worden getraind met Proximal Policy Optimization (PPO).

• Stuurcoëfficiënt ($c$): Wordt adaptief bepaald door het gemiddelde van activatiegroottes van correct voorspelde trainingssteken, waardoor handmatige afstemming over lagen overbodig wordt.
• Criticusnetwerk: Een 2-laags MLP dat de waardenfunctie $V_\phi(s)$ schat, waardoor het mogelijk is beleidsbeperkingen te scheiden van fouten in waardeschatting tijdens de analyse.

3. Belangrijkste Bijdragen

Het artikel schetst vier primaire bijdragen:

1. Token-niveau Interpreteerbare Interventielogboeken: Een raamwerk dat logboeken produceert die exact identificeren welk SAE-kenmerk bij elke token werd geselecteerd en wat het effect was op de output, waardoor gedetailleerde gedragsanalyse mogelijk wordt.
2. Diagnostische Hulpmiddelen:
   • Branch Point Tracking: Identificeert kritieke beslispunten waar identieke voorvoegsels verschillende uitkomsten opleveren op basis van kenmerkselectie.
   • Criticus Trajectanalyse: Gebruikt de geleerde waardenfunctie om onderscheid te maken tussen beleidsfalen (kenmerkselectie) en criticusfalen (beloningsschatting).
3. Adaptieve Kenmerkmaskering (AFM): Een mechanisme dat exploratie en exploitatie in evenwicht brengt om diverse kenmerkontdekking te waarborgen zonder de interpreteerbaarheid van interventies met één kenmerk op te offeren.
4. Empirische Karakterisering van Kenmerksemantiek: Toont aan dat vroege lagen syntactische structuur coderen terwijl latere lagen abstracte semantiek vastleggen, en dat verschillende taken onderscheidende bottleneckpatronen vertonen.

4. Experimentele Resultaten

Het raamwerk werd geëvalueerd op Gemma 2 2B (en LLaMA-3.1 8B) met behulp van Gemma Scope en LLaMA Scope SAE's over vijf benchmarks: MMLU, BBQ, GSM8K, HarmBench en XSTest.

• Prestatiewinst: CRL boekte verbeteringen op alle benchmarks. Opvallende winsten zijn:
  • HarmBench: +7,66 procentpunten (41,46% $\to$ 49,12%).
  • BBQ (Ambigu): +5,69 procentpunten.
  • MMLU: +3,31 procentpunten.
  • GSM8K: +1,03 procentpunten.
• Complementariteit: CRL levert additieve winst wanneer toegepast bovenop Supervised Fine-Tuned (SFT) modellen (bijvoorbeeld steeg de MMLU-accuraatheid van 55,2% naar 56,1% met SFT + CRL).
• Laagspecifieke Semantiek:
  • Laag 10: Kenmerken vangen concrete, syntactische eigenschappen (bijvoorbeeld "wiskundige notatie", "voornaamwoordverwijzingen").
  • Laag 20: Kenmerken vangen abstracte, contextuele semantiek (bijvoorbeeld "logische afleidingsstructuur", "wiskundige bewijzen").
  • De optimale interventielagen variëren per taak (bijvoorbeeld Laag 5 voor BBQ, Laag 24 voor MMLU).
• Criticusanalyse:
  • BBQ: Toont effectieve coördinatie tussen beleid en criticus met duidelijke discriminatie van stekels.
  • MMLU: Vertoont een "criticus bottleneck" waarbij gecorrigeerde en misleidende stekels bijna niet van elkaar te onderscheiden zijn, wat suggereert dat de criticus moeite heeft met het schatten van waarde voor complexe redeneertaken.
  • HarmBench versus XSTest: HarmBench toont duidelijke gaten tussen correcte/incorrecte stekels, terwijl XSTest omgekeerde schattingspatronen vertoont, wat wijst op taakafhankelijke uitdagingen in beloningsschatting.

5. Betekenis en Claims

Het artikel positioneert CRL als een significante vooruitgang in mechanische interpretatie. De primaire betekenis ligt in de verschuiving van het paradigma van statische kenmerkanalyse (observeren wat activeert) naar dynamische interventieprobes (observeren wat het model verandert).

• Interpreteerbaarheid: Door selectie van één kenmerk af te dwingen en interventies te loggen, biedt CRL een "mechanisch" beeld van modelgedrag, wat aantoont dat kenmerken zoals "wiskundige bewerkingen" lexicaal kunnen generaliseren naar kwantitatieve eenheden zoals "baby's" in woordproblemen.
• Diagnostisch Vermogen: Het raamwerk dient als diagnostisch hulpmiddel voor zowel het model als de SAE. Als een beleid instort tot semantisch oncoherente kenmerken ondanks AFM, wijst dit op een tekortkoming in de SAE (bijvoorbeeld ontbrekende concepten of polysemantiteit) in plaats van een falen van de controlemethode.
• Praktijktoepasbaarheid: CRL werkt met een lage inferentie-overhead (8–9% latentie) en vereist geen hertraining van de basismodelgewichten, wat een lichtgewicht, post-hoc verbeteringsmechanisme biedt.

De auteurs concluderen dat geleerde kenmerksturing, wanneer gekoppeld aan reinforcement learning, een praktisch hulpmiddel vestigt voor het begrijpen en controleren van taalmodellen, waarbij statische analyse wordt aangevuld met dynamische, taakspecifieke inzichten.