CorrSteer: Generation-Time LLM Steering via Correlated Sparse Autoencoder Features

TL;DR

CorrSteer tackles the challenge of steering LLM behavior with Sparse Autoencoders (SAEs) without requiring large contrastive datasets or heavy activation storage. It introduces a generation-time feature-selection pipeline based on Pearson correlation between SAE activations and task outcomes, deriving steering coefficients from positive outcomes and applying additive steering vectors during inference. The method yields improvements across QA, bias mitigation, safety, and reasoning benchmarks on Gemma-2 2B and LLaMA-3.1 8B, including notable gains such as +3.3% in MMLU with 4000 samples and +27.2% in HarmBench with only 108 samples, while offering low side effects via the Side Effect Ratio (SER). By revealing task-aligned circuits and maintaining interpretability through SAE features, CorrSteer demonstrates a scalable, automated approach to generation-time LLM steering using correlation-based feature selection.

Abstract

Sparse Autoencoders (SAEs) can extract interpretable features from large language models (LLMs) without supervision. However, their effectiveness in downstream steering tasks is limited by the requirement for contrastive datasets or large activation storage. To address these limitations, we propose CorrSteer, which selects features by correlating sample correctness with SAE activations from generated tokens at inference time. This approach uses only inference-time activations to extract more relevant features, thereby reducing spurious correlations. It also obtains steering coefficients from average activations, automating the entire pipeline. Our method shows improved task performance on QA, bias mitigation, jailbreaking prevention, and reasoning benchmarks on Gemma-2 2B and LLaMA-3.1 8B, notably achieving a +3.3% improvement in MMLU performance with 4000 samples and a +27.2% improvement in HarmBench with only 108 samples. Selected features demonstrate semantically meaningful patterns aligned with each task's requirements, revealing the underlying capabilities that drive performance. Our work establishes correlation-based selection as an effective and scalable approach for automated SAE steering across language model applications.

Figures (35)

• Figure 1: System diagram of CorrSteer. CorrSteer selects task-relevant SAE features by correlating generated-token activations with outcomes, and constructs steering vectors applied as CorrSteer-S, CorrSteer-A, or CorrSteer-P. This generation-time steering shifts model behavior from unintended to intended responses while reducing side effects.
• Figure 2: Comparison of features selected by CorrSteer-S, CorrSteer-A, and CorrSteer-P on BBQ (disambiguous) across all Gemma-2 2B layers. Red points denote selected features.
• Figure 3: Relation between sample counts and test performance, final matched count of selected features, and most correlated features from each Gemma-2 2B layer. Dotted lines show baseline default LLM performance and constrained decoding performance on MMLU answer options.
• Figure 4: SER comparison between different CorrSteer variants for Gemma-2 2B.
• Figure 5: Benchmark performance of CorrSteer variants compared with the non-steered model on LLaMA-3.1 8B.