💬 NLP

Control Reinforcement Learning: Interpretable Token-Level Steering of LLMs via Sparse Autoencoder Features

This paper introduces Control Reinforcement Learning (CRL), a framework that trains a policy to dynamically select and amplify interpretable sparse autoencoder features at the token level to steer LLM outputs, thereby providing per-token intervention protocols and new mechanistic insights into model behavior across various benchmarks.

Original authors: Seonglae Cho, Zekun Wu, Adriano Koshiyama

Published 2026-05-05

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Seonglae Cho, Zekun Wu, Adriano Koshiyama

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a large language model (like the one powering this chat) as a vast, complex orchestra. In this orchestra, thousands of musicians (neurons) play simultaneously. Sometimes they perform a beautiful symphony, but at other times they become confused and play the wrong notes, leading to hallucinations or poor responses.

For a long time, scientists have tried to understand this orchestra by listening to the music and guessing which instruments are playing. They discovered "Sparse Autoencoders" (SAEs), which act like an all-powerful sound engineer. The SAE can decompose the music into individual, named tracks: "Math Track," "Safety Track," "Grammar Track," and so on.

The Problem:
Knowing which tracks are playing is not enough. Just because the "Math Track" is active does not mean that turning up the volume on this track will solve a math problem. Sometimes, turning up a track makes things even worse. Previous methods could tell you what was played, but not what would happen if you changed the volume.

The Solution: Control Reinforcement Learning (CRL)
This work introduces a new method called Control Reinforcement Learning (CRL). Imagine CRL as the tuning of a clever conductor who stands directly beside the orchestra during the performance.

Here is how it works, step by step:

1. The Conductor's Task (The Policy)

For every single word (token) the model is about to say, this conductor examines the current state of the orchestra. The conductor has a remote control with buttons for each individual "track" (feature) identified by the SAE.

The Decision: The conductor decides: "Right now we need to turn up the 'Logical Reasoning' track" or "Let's turn down the 'Emotional Bias' track."
The Action: The conductor immediately adjusts the volume of that specific track for just that one word.

2. Learning by Doing (Reinforcement Learning)

The conductor does not know the rules initially. He learns by playing the game:

If the orchestra plays a correct answer, the conductor receives a "reward" (a point).
If the orchestra plays an incorrect answer, the conductor receives no points.
Over time, the conductor learns exactly which tracks need to be amplified at which moments to achieve the best results.

3. The "Adaptive Mask" (Preventing Bad Habits)

A clever conductor might become lazy and simply turn up the same "Math Track" for every problem. To prevent this, the work uses a trick called Adaptive Feature Masking.

Imagine the conductor has a rule: "You may not use the same instrument twice in a row unless you have tried others first."
This forces the conductor to explore different parts of the orchestra and discover new ways to solve problems, rather than relying on a single trick.

4. The "Logbook" (Interpretability)

This is the greatest breakthrough of the work. Since the conductor makes a specific choice for every single word, we obtain a detailed logbook of the entire performance.

We can look back and say: "Ah, at word number 5, the conductor turned up the 'Safety Track,' which prevented the model from saying something rude."
We can see exactly why the model succeeded or failed. It is as if we have a transcript of the conductor's thoughts for every note played.

What Did They Discover?

By applying this method to a model named Gemma 2, the researchers discovered some fascinating things:

Early vs. Late Layers: The orchestra has different sections. The "early" sections (layers) handle grammar and structure (such as punctuation and sentence construction). The "late" sections handle deep meaning and logic. The work found that to solve a math problem, one often needs to adjust the "late" sections where logic resides, not just the "early" sections where words are formed.
The "Branching Points": Sometimes two very similar questions require completely different conductor actions to obtain the correct answer. The system can identify the exact moment when the path between a correct and an incorrect answer diverges.
Safety vs. Bias: The system learned that fixing "bias" (injustice) and fixing "safety" (rejecting harmful requests) require different strategies. For some tasks, the conductor must be very specific; for others, he must explore many different tracks.

The Results

The work tested this on various challenges:

Math (GSM8K): The conductor helped the model solve more math problems by amplifying the correct logical tracks.
Safety (HarmBench): The model became better at refusing to do harmful things without being overly cautious.
Knowledge (MMLU): The model provided more accurate answers to knowledge questions.

Summary:
This work does not just tell us what the AI thinks; it gives us a tool to control the AI's thinking in real time, word by word. It transforms a "black box" into a transparent machine where we can see exactly which internal switches were flipped to obtain the correct answer. It is like giving a human a remote control to steer the inner thoughts of the AI and ensure it stays on the right path.

Technical Summary: Control Reinforcement Learning (CRL)

1. Problem Statement

Sparse Autoencoders (SAEs) have proven to be a powerful tool for decomposing activations in language models into interpretable, human-readable features. However, existing analysis methods are largely static: they identify which features activate for specific inputs but do not reveal which features, when reinforced, actually change the model outputs. Current activation manipulation techniques often rely on contrastive datasets for feature selection or extensive storage of activations to optimize coefficients, lacking adaptive feedback mechanisms to adjust control strategies based on generation quality.

This work bridges the gap between static feature identification and dynamic, interpretable intervention. The core challenge is to develop a framework that not only selects SAE features to control a model but also provides token-level intervention protocols explaining why a feature was selected and how it influenced the output, thereby transforming control into a tool for mechanistic interpretability.

2. Methodology: Control Reinforcement Learning (CRL)

The authors propose Control Reinforcement Learning (CRL), a framework that treats the control of transformer representations as a Markov Decision Process (MDP) over SAE features.

2.1. MDP Formulation

State ( $S$ ): The observation is the activation of the residual stream ( $x \in \mathbb{R}^d$ ) in a specific layer $\ell$ for the current token position.
Action ( $A$ ): A binary vector selecting one or more SAE features for reinforcement. The policy selects the top- $k$ features (typically $k=1$ for interpretability).
Transition ( $P$ ): Deterministic, governed by the forward pass of the transformer with the applied feature intervention.
Reward ( $R$ ): Task-specific rewards based on the quality of the final output (e.g., correctness on MMLU, refusal on safety tasks).

2.2. CRL Token and Feature Control

At each generation step $t$ , the framework performs token-level interventions:
$\tilde{x}_t = x_t + c \cdot a_t W_{dec}$
Here, $a_t$ is the sparse binary selection vector, $c$ is the control coefficient, and $W_{dec}$ is the SAE decoder. The policy network $\pi_\theta$ maps residual stream activations to logits for feature selection and is implemented as a 2-layer MLP.

2.3. Adaptive Feature Masking (AFM)

To prevent the policy from collapsing into a small set of frequently active features, CRL introduces Adaptive Feature Masking.

A per-sample Boolean mask restricts feature selection options.
The mask begins with a small subset of features and progressively expands during generation, integrating features that were naturally activated in previous steps.
This mechanism promotes the discovery of diverse features while preserving the interpretability of single-feature interventions.

2.4. Optimization

The policy and a critic network (for state value estimation) are trained using Proximal Policy Optimization (PPO).

Control Coefficient ( $c$ ): Adaptively determined by averaging activation strengths from correctly predicted training samples, eliminating the need for manual tuning across layers.
Critic Network: A 2-layer MLP that estimates the value function $V_\phi(s)$ , enabling the separation of policy errors (feature selection) from critic errors (reward estimation) during analysis.

3. Main Contributions

The work outlines four primary contributions:

Token-Level Intervention Protocols: A framework that generates protocols precisely identifying which SAE feature was selected at each token and what effect it had on the output, enabling fine-grained behavioral analysis.
Diagnostic Tools:
- Branch Point Tracking: Identifies critical decision points where identical prefixes lead to different outcomes based on feature selection.
- Critic Trajectory Analysis: Uses the learned value function to distinguish between policy errors (feature selection) and critic errors (reward estimation).
Adaptive Feature Masking (AFM): A mechanism balancing exploration and exploitation to ensure diverse feature discovery without sacrificing the interpretability of single-feature interventions.
Empirical Characterization of Feature Semantics: Demonstrates that early layers encode syntactic structures while later layers capture abstract semantics, and that different tasks exhibit distinct bottleneck patterns.

4. Experimental Results

The framework was evaluated on Gemma 2 2B (and LLaMA-3.1 8B) using Gemma Scope and LLaMA Scope SAEs across five benchmarks: MMLU, BBQ, GSM8K, HarmBench, and XSTest.

Performance Improvements: CRL achieved improvements across all benchmarks. Notable gains include:
- HarmBench: +7.66 percentage points (41.46% $\to$ 49.12%).
- BBQ (Ambiguous): +5.69 percentage points.
- MMLU: +3.31 percentage points.
- GSM8K: +1.03 percentage points.
Complementarity: CRL yields additive gains when applied to Supervised Fine-Tuned (SFT) models (e.g., MMLU accuracy increased from 55.2% to 56.1% with SFT + CRL).
Layer-Specific Semantics:
- Layer 10: Features capture concrete, syntactic properties (e.g., "mathematical notation", "pronoun references").
- Layer 20: Features capture abstract, contextual semantics (e.g., "structure of logical derivations", "mathematical proofs").
- Optimal intervention layers vary by task (e.g., Layer 5 for BBQ, Layer 24 for MMLU).
Critic Analysis:
- BBQ: Shows effective policy-critic coordination with clear sample discrimination.
- MMLU: Exhibits a "critic bottleneck" where corrected and misdirected samples are barely distinguishable, suggesting the critic struggles to estimate value for complex reasoning tasks.
- HarmBench vs. XSTest: HarmBench shows clear gaps between correct and incorrect samples, whereas XSTest displays inverted estimation patterns, indicating task-specific challenges in reward estimation.

5. Significance and Claims

The work positions CRL as a significant advancement in mechanistic interpretability. Its primary significance lies in the paradigm shift from static feature analysis (observing what activates) to dynamic intervention sampling (observing what changes the model).

Interpretability: By enforcing the selection of single features and logging interventions, CRL provides a "mechanistic" view of model behavior, demonstrating that features like "mathematical operations" can lexically generalize to quantitative units like "babies" in word problems.
Diagnostic Capability: The framework serves as a diagnostic tool for both the model and the SAE. If a policy collapses onto semantically incoherent features despite AFM, it signals an SAE deficiency (e.g., missing concepts or polysemy) rather than a failure of the control method.
Practicality: CRL operates with low inference overhead (8–9% latency) and requires no retraining of the model's base weights, offering a lightweight, post-hoc improvement mechanism.

The authors conclude that learning to steer features in combination with reinforcement learning establishes a practical tool for understanding and controlling language models, complementing static analysis with dynamic, task-specific insights.