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

TL;DR

This paper introduces Control Reinforcement Learning (CRL), a method for interpretable, token-level steering of language models using sparse autoencoder features, enabling dynamic analysis and improved performance.

Contribution

CRL is a novel framework that trains policies to select features for model steering, providing dynamic interpretability and analysis tools for language models.

Findings

CRL achieves performance improvements on multiple benchmarks.

It provides interpretable, per-token intervention logs.

Layer-wise analysis reveals syntactic and semantic features.

Abstract

Sparse autoencoders (SAEs) decompose language model activations into interpretable features, but existing methods reveal only which features activate, not which change model outputs when amplified. We introduce Control Reinforcement Learning (CRL), which trains a policy to select SAE features for steering at each token, producing interpretable intervention logs: the learned policy identifies features that change model outputs when amplified. Adaptive Feature Masking encourages diverse feature discovery while preserving singlefeature 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; layer-wise comparison reveals syntactic features in early layers and semantic features in later layers. On Gemma 2 2B across MMLU, BBQ, GSM8K, HarmBench, and XSTest, CRL achieves improvements while providing per-token intervention logs. These results establish learned feature steering as a mechanistic interpretability tool that complements static feature analysis with dynamic intervention probes