app/src/content/article.mdx

---
seoThumbImage: 'https://seongland.com/article/crl/og.png'
title: 'Control Reinforcement Learning'
subtitle: >-
  Interpretable Token-Level Steering of LLMs via SAE Features
description: 'CRL: Interpretable Token-Level Steering of LLMs via Sparse Autoencoder Features'
authors:
  - name: Seonglae Cho
    affiliations:
      - 1
      - 2
  - name: Zekun Wu
    affiliations:
      - 1
      - 2
  - name: Adriano Koshiyama
    affiliations:
      - 1
      - 2
affiliations:
  - name: Holistic AI
    url: 'https://holisticai.com'
  - name: University College London
    url: 'https://ucl.ac.uk'
published: 'Feb. 13, 2026'
licence: >
  Diagrams and text are licensed under <a
  href="https://creativecommons.org/licenses/by/4.0/" target="_blank"
  rel="noopener noreferrer">CC‑BY 4.0</a>.
doi: '10.48550/arXiv.2602.10437'
arxiv: '2602.10437'
tags:
  - mechanistic interpretability
  - sparse autoencoders
  - reinforcement learning
  - language models
  - feature attribution

tableOfContentsAutoCollapse: true
---

import Sidenote from '../components/Sidenote.astro';
import HtmlEmbed from '../components/HtmlEmbed.astro';
import Note from '../components/Note.astro';
import Wide from '../components/Wide.astro';
import FullWidth from '../components/FullWidth.astro';
import Quote from '../components/Quote.astro';
import Reference from '../components/Reference.astro';

## Abstract

Watching a feature light up tells you it was present. It does not tell you it did anything. [Sparse autoencoders](https://texonom.com/1fa3c713076d4e34aa43f7c1dce62a94) (SAEs) decompose language model activations into interpretable features [@bricken2023monosemanticity; @huben2024sparse], and almost everything built on them reads off *which* features activate, not which ones actually **change model outputs** when amplified. We introduce **Control [Reinforcement Learning](https://texonom.com/3e6e1e5b806941b6b586ce2c3d1bbd02) (CRL)**, which trains a policy to select SAE features for steering at each token, producing interpretable intervention logs.

**Adaptive Feature Masking** encourages diverse feature discovery while preserving single-feature interpretability. The framework yields new analysis capabilities: **branch point tracking** locates tokens where feature choice determines output correctness; **critic trajectory analysis** separates policy limitations from value estimation errors; and **layer-wise comparison** reveals syntactic features in early layers and semantic features in later layers [@tenney-etal-2019-bert].

On [Gemma 2](https://texonom.com/8aa07c8eba2c4d4ba61213b9a4f711c5) 2B [@gemmateam2024gemma2improvingopen] across [MMLU](https://texonom.com/d2a91b61ee71491d9e6e284a5bdcc709) [@hendrycks2021measuring], [BBQ](https://texonom.com/14bc3c96247d81038bb3c4aa92f78a0d) [@parrish-etal-2022-bbq], GSM8K [@cobbe2021trainingverifierssolvemath], HarmBench [@mazeika2024harmbench], and [XSTest](https://texonom.com/237c3c96247d8090a166c4cbee0dc287) [@rottger-etal-2024-xstest], CRL achieves consistent improvements while providing per-token intervention logs. These results establish learned feature steering as a [mechanistic interpretability](https://texonom.com/e40305fd878c46ca85d99ea93ee9a2ff) tool that complements static feature analysis with dynamic intervention probes.

<Wide>
<HtmlEmbed src="steering-demo.html" title="Token-Level Steering Simulation" desc="Figure 2: Watch how different SAE features steer token generation along divergent paths. The same prompt context leads to correct or incorrect answers depending on which feature is amplified. Feature IDs link to Neuronpedia for detailed activation analysis." />
</Wide>

## The Problem: Static Feature Analysis Falls Short

Sparse autoencoders have emerged as a powerful tool for decomposing the internal representations of language models. By training an encoder-decoder pair with a sparsity penalty, SAEs recover interpretable features from the [residual stream](https://texonom.com/beb7f1e3c9334362b518f0350300780a) — components that correspond to concepts like "mathematical reasoning" or "safety refusal patterns."

<Sidenote>
SAEs learn: $\mathcal{L} = \|\mathbf{x} - \hat{\mathbf{x}}\|^2 + \lambda \|\mathbf{z}\|_1$ where $\mathbf{z}$ is the sparse feature activation vector.
</Sidenote>

But there's a fundamental limitation: **static analysis tells you which features are present, not which ones matter.** Knowing that [Feature 7708](https://neuronpedia.org/gemma-2-2b/20-gemmascope-res-16k/7708) ("mathematical operations") activates on a math problem doesn't tell you whether *amplifying* that feature would improve the model's answer, make it worse, or have no effect at all.

This gap between activation analysis and causal intervention is where CRL enters. Instead of cataloging which features fire, we ask: **which features, when amplified, actually change the model's output?**

<Note variant="info">
**The key distinction:** Traditional SAE analysis is like reading a map — you can see the terrain. CRL is like walking through it — you discover which paths actually lead somewhere different.
</Note>

Current approaches to SAE-based control face limitations:
- **Contrastive methods** require carefully curated positive/negative datasets for each feature
- **[Activation patching](https://texonom.com/11fc3c96247d80d683e6c67f149c992d)** requires extensive storage and computation for exhaustive feature probing
- **Manual coefficient tuning** doesn't generalize across layers or tasks

CRL addresses these by framing feature selection as a *learned* decision at every token, using reinforcement learning to discover which features are worth amplifying.


## How CRL Works

<Wide>
<HtmlEmbed src="crl-flowchart-animated.html" title="CRL System Overview" desc="Figure 3: Interactive CRL system diagram. Hover over components to see their role. The policy network observes the residual stream and selects SAE features for steering, while the critic estimates expected outcomes for PPO training." />
</Wide>

### Formulating Steering as an MDP

CRL treats the control of transformer representations as a **[Markov Decision Process](https://texonom.com/d01c5c23d56648ce843727514008e25c)** over SAE features:

<div class="crl-equation">

$$\tilde{\mathbf{x}}_t = \mathbf{x}_t + c \cdot \mathbf{a}_t \mathbf{W}_{dec}$$

</div>

where $\mathbf{x}_t$ is the residual stream activation at token $t$, $\mathbf{a}_t$ is the binary feature selection vector, $c$ is the steering coefficient, and $\mathbf{W}_{dec}$ is the SAE decoder matrix.

The MDP components are:

- **State** $\mathbf{s} = \mathbf{x}$: the residual stream activation at the intervention layer
- **Action** $\mathbf{a} \in \{0,1\}^{d_{dict}}$: a sparse binary vector selecting SAE features (typically $k=1$)
- **Transition**: deterministic, governed by the transformer's forward pass with the steering intervention
- **Reward**: task-specific binary signal (correct/incorrect output)

<Sidenote>
Using $k=1$ (single feature selection) is critical for interpretability — each intervention is attributable to exactly one SAE feature.
</Sidenote>

### The Policy and Critic Networks

<Wide>
<HtmlEmbed src="policy-critic-diagram.html" title="Policy & Critic Architecture" desc="Figure: The policy observes the residual stream per token and selects SAE features for steering. The critic estimates state values for advantage computation from verifiable rewards." />
</Wide>

The **policy network** $\pi_\theta$ is a 2-layer MLP with Tanh activation that maps residual stream activations to feature selection logits:

$$\boldsymbol{\mu} = \pi_\theta(\mathbf{s}), \quad \mathbf{a} = \text{TopK}(\boldsymbol{\mu}, k)$$

During training, actions are sampled from a categorical distribution over softmax probabilities for unbiased [policy gradients](https://texonom.com/331be391bdf44796a94ac9d6ddfe6c01). During evaluation, we use deterministic argmax for reproducibility.

The **critic network** $V_\phi$ estimates the state [value function](https://texonom.com/d9d3467fea994443b2336fd7958dc00b), enabling advantage estimation for [PPO](https://texonom.com/87ce8ebe81f84ac1a7617b1f6def9e26) [@schulman2017proximal] training. Crucially, the critic also serves as an *analysis tool* — its learned value estimates reveal where the policy succeeds and fails.

<Note variant="info">
**Why PPO over [DPO](https://texonom.com/d4099259ad6d443e8f9c2c9d1b45d6db)/[GRPO](https://texonom.com/187c3c96247d80dfb88bceaf4abad43c)?** [@shao2024deepseekmathpushinglimitsmathematical] PPO provides an explicit value function that we analyze as a diagnostic tool. The critic trajectories reveal bottleneck patterns that would be invisible with critic-free methods.
</Note>

### Partial Observability

CRL-Token operates under partial observability: the policy sees only the residual stream at layer $\ell$, without access to the tokens that will be sampled. We mitigate this by using temperature-0 sampling, making the generation deterministic given the residual stream state.


## Adaptive Feature Masking

A naive policy would collapse to repeatedly selecting the same high-reward feature. **Adaptive Feature Masking (AFM)** prevents this by dynamically managing which features are available for selection.

<div class="crl-equation">

$$\text{mask}_i^{(t+1)} = \text{mask}_i^{(t)} \lor \text{active}_i^{(t)}$$

</div>

The mask starts small and **progressively expands** as the model generates tokens — features that naturally activate during generation become available for future policy selections:

1. Policy computes logits and masks unavailable features: $\mathbf{a}_t = \text{TopK}(\boldsymbol{\mu} \odot \text{mask}^{(t)}, k)$
2. Steering is applied to the residual stream
3. Natural SAE activations at the current token expand the mask for the next step

This creates a virtuous cycle: early steering may activate new features that become candidates for later steering, encouraging diverse feature discovery while building on the model's own representations.

<Note variant="success">
AFM is particularly important for multi-token tasks like GSM8K, where the reasoning chain spans hundreds of tokens and different features become relevant at different stages.
</Note>


## Branch Analysis: Where Steering Decides the Answer

**Branch analysis** identifies tokens where the *choice of feature* determines whether the model produces a correct or incorrect answer.

### How Branch Analysis Works

For each generation, we track:
1. The **context** (shared prefix before the branch point)
2. The **correct path** — the feature that leads to a correct answer
3. The **incorrect path** — the feature that leads to a wrong answer

Because the two paths share the prefix, the branch point localizes where the choice of feature is causally responsible for the outcome.

### Layer-Specific Feature Semantics

Comparing features selected at Layer 10 vs Layer 20 reveals a striking pattern:

<div class="crl-comparison">
  <div class="crl-comparison__side crl-comparison__side--correct">
    <h4>Layer 10: Syntactic Features</h4>
    <ul>
      <li><a href="https://neuronpedia.org/gemma-2-2b/10-gemmascope-res-16k/16145"><span class="crl-layer-badge">L10/16145</span></a> "mathematical notation"</li>
      <li><a href="https://neuronpedia.org/gemma-2-2b/10-gemmascope-res-16k/5423"><span class="crl-layer-badge">L10/5423</span></a> "pronoun references"</li>
      <li><a href="https://neuronpedia.org/gemma-2-2b/10-gemmascope-res-16k/8428"><span class="crl-layer-badge">L10/8428</span></a> "clinical examinations" (surface similarity)</li>
      <li><a href="https://neuronpedia.org/gemma-2-2b/10-gemmascope-res-16k/8775"><span class="crl-layer-badge">L10/8775</span></a> "legal contracts" (terminology match)</li>
      <li><a href="https://neuronpedia.org/gemma-2-2b/10-gemmascope-res-16k/12607"><span class="crl-layer-badge">L10/12607</span></a> "mathematical proofs" (logical structure)</li>
    </ul>
  </div>
  <div class="crl-comparison__side crl-comparison__side--incorrect">
    <h4>Layer 20: Semantic Features</h4>
    <ul>
      <li><a href="https://neuronpedia.org/gemma-2-2b/20-gemmascope-res-16k/11163"><span class="crl-layer-badge">L20/11163</span></a> "assessments and strategic plans"</li>
      <li><a href="https://neuronpedia.org/gemma-2-2b/20-gemmascope-res-16k/2878"><span class="crl-layer-badge">L20/2878</span></a> "code formatting" (structural patterns)</li>
      <li><a href="https://neuronpedia.org/gemma-2-2b/20-gemmascope-res-16k/998"><span class="crl-layer-badge">L20/998</span></a> "needs identification" (requirement semantics)</li>
    </ul>
  </div>
</div>

Layer 10 features capture **concrete syntactic patterns** like surface-level notation and entity tracking. Layer 20 features encode **abstract semantics** like planning and strategic reasoning.

<Note variant="danger">
**Feature labels have limited predictive value.** A feature labeled "emotional relationships" can solve speed calculations, while "mathematical notation" fails on some arithmetic. So learned steering effects diverge from semantic interpretations derived from activation patterns alone.
</Note>

Out of 8 analyzed branch pairs, Layer 10 features succeed on **3 cases** (discrete counting, direct arithmetic), while Layer 20 features succeed on **5 cases** (multi-step accumulation, variable reasoning).


## Empirical Validation

To confirm that the learned policy selects *meaningful* features — not arbitrary ones — we compare CRL against baselines on seven benchmarks. The goal is not performance gain per se, but to validate that the policy has learned a non-trivial feature selection strategy worth analyzing.

<HtmlEmbed src="results-table.html" title="CRL Steering Results" desc="Table 1: Before vs. after CRL steering. Toggle between Gemma 2 2B and [LLaMA](https://texonom.com/f2b6721202d44d469add84d8a366809c) 3.1 8B to compare across tasks and architectures." />

### Learned vs. Random Selection

<Wide>
<HtmlEmbed src="baseline-comparison.html" title="Steering Baseline Comparison" desc="Figure 5: Random feature selection can degrade performance (−1.81 on BBQ), confirming the policy learns non-trivial feature selection." />
</Wide>

| Method | BBQ Ambig | HarmBench |
|--------|-----------|-----------|
| No Steering (Base) | 60.17 | 41.46 |
| Random Feature | 58.36 | 45.35 |
| Random + AFM | 60.16 | 46.96 |
| Most-Active Feature | 59.94 | 48.03 |
| **CRL (Ours)** | **65.86** | **49.12** |

Random feature steering can actually *hurt* performance (−1.81 on BBQ Ambig), while selecting the most-active SAE feature slightly hurts on BBQ Ambig. The learned policy discovers features whose *causal effects* differ from their activation magnitude — precisely the gap that static analysis cannot reveal.


## Critic Trajectory Analysis

The learned critic (value function) reveals *where* in the generation process the expected reward changes — and whether task limitations stem from the policy or the value estimator.

<Wide>
<HtmlEmbed src="critic-trajectories.html" title="Interactive Critic Trajectory Explorer" desc="Figure 7: Explore critic value trajectories across tasks. GSM8K shows multi-token generation (~1000 steps); HarmBench and XSTest show short-form generation (~30 steps). Toggle between tasks to compare." />
</Wide>

### Key Observations

For multi-token generation (GSM8K), critic trajectories reveal the *temporal dynamics* of steering:

- **Correct samples** (green): Higher gradients with divergence around token 200–400
- **Incorrect samples** (red): Rising values below the correct samples
- **Corrected samples** (blue): Show value increases at the ~400-token point, indicating interventions become more effective in later reasoning stages
- **Misguided samples** (yellow): Rising values that drop late — the steering seemed promising but ultimately led astray


## Feature Interpretability

CRL provides per-token intervention logs that reveal *which* features are selected and *when*. This enables a new form of mechanistic analysis.

### Feature-Token Alignment

<Wide>
<HtmlEmbed src="token-highlights.html" title="Token-Level Feature Activation Highlights" desc="Figures 8-11: Interactive token-level visualization of SAE feature activations on GSM8K reasoning chains. Corrected examples show feature activations aligning with mathematical tokens; misguided examples show token-feature misalignment. Hover tokens for details; click feature pills to view on Neuronpedia." />
</Wide>

**Corrected examples** show clear semantic coherence: features activate on mathematical tokens, operators, and quantitative units -- demonstrating lexical generalization beyond surface forms.

### When Features Fail

**Misguided examples** (select the "Misguided" tab above) reveal failure modes:
- Task-irrelevant features disrupting the reasoning chain
- Token-feature misalignment producing incoherent activations
- Feature interference causing the model to abandon its solution path

### Feature Diversity

<HtmlEmbed src="feature-frequency.html" title="Feature Selection Frequency" desc="Figure 13: Feature selection frequency follows a rank-dependent decay, confirming the policy learns selective (not random) feature choices. Click feature IDs to view on Neuronpedia." />

Feature diversity varies significantly across tasks, measured by entropy of the feature selection distribution:

| Task | Feature Diversity | Interpretation |
|------|------------------|----------------|
| GSM8K | 6.695 | Complex reasoning needs diverse features |
| MMLU-Pro | 5.476 | High complexity, broad engagement |
| XSTest | 4.757 | Intermediate diversity |
| HarmBench | 2.938 | Focused refusal uses concentrated features |

Complex reasoning tasks require *broader* feature engagement, while focused behavioral tasks benefit from *concentrated* feature activation.


## Layer-Wise Analysis

Different layers encode different types of information, and CRL's layer-wise training reveals this structure.

<Wide>
<HtmlEmbed src="layer-heatmap.html" title="Layer × Coefficient Sensitivity" desc="Figure 14: Interactive heatmap showing accuracy across layer-coefficient combinations. Each task has a different optimal configuration. Hover cells for details." />
</Wide>

### Residual Stream Norms

<HtmlEmbed src="norm-growth.html" title="Residual Stream Norm Growth" desc="Figure 15: Residual stream norms increase across layers, requiring larger steering coefficients in deeper layers. Hover to see exact values and coefficient ranges." />

Residual stream norms increase with depth, which directly impacts the optimal steering coefficient:
- **Early layers** (L5–L10): $c = 10\text{–}20$
- **Middle layers** (L15–L20): $c = 30\text{–}60$
- **Late layers** (L22–L25): $c = 60\text{–}100$

Large coefficients degrade early-layer performance, while late layers tolerate — and require — stronger interventions.


## Discussion

### What CRL Reveals as an Analysis Tool

CRL provides three diagnostic capabilities that complement static feature analysis:

**1. Branch tracking** locates critical decision points — tokens where the choice of a single SAE feature determines whether the output is correct or incorrect. This localizes where the feature choice is causally responsible, which correlational activation analysis cannot show.

**2. Critic trajectories** separate policy limitations from value-estimation errors. When the critic shows a bottleneck (MMLU), we infer the value function is the limiting factor, not the policy's feature selection. When the critic shows clear discrimination (BBQ), the policy-critic coordination is effective.

**3. Layer comparison** reveals a syntactic-to-semantic gradient: early layers encode concrete patterns (notation, entity tracking), while later layers capture abstract structure (strategic planning).

### Limitations

- **Layer selection is post-hoc:** optimal layers are identified via validation, not predicted a priori
- **Feature labels are unreliable:** SAE feature descriptions often diverge from their actual steering effects
- **Single-feature constraint:** selecting $k=1$ maximizes interpretability but may leave performance on the table
- **Computational cost:** while modest (8–9% overhead), CRL requires PPO training for each task-layer combination


## Conclusion

Control Reinforcement Learning transforms SAE feature steering from a static analysis tool into a **dynamic, token-level intervention framework** with built-in interpretability. By training a policy to select features at each generation step, CRL produces mechanistic analysis tools:

- **Branch tracking** locates critical decision points where feature choice determines output correctness
- **Critic trajectories** separate policy limitations from value-estimation errors
- **Layer comparison** reveals syntactic-to-semantic feature gradients across network depth

Validated across knowledge, bias, reasoning, and safety benchmarks, CRL maintains single-feature interpretability throughout. The framework establishes learned feature steering as a practical mechanistic interpretability tool that complements static feature analysis with dynamic intervention probes.