CorrSteer: Steering Improves Task Performance and Safety in LLMs through Correlation-based Sparse Autoencoder Feature Selection

2025-08-17

AI Summary

Sparse autoencoders (SAEs) enable unsupervised extraction of interpretable features from large language models (LLMs), but existing steering methods rely either on contrastive datasets or massive activation caches—limiting practical deployment. This paper introduces CorrSteer, the first method to leverage the correlation between sample correctness and SAE activations *during inference* to automatically identify semantically meaningful, non-spurious features and jointly learn steering coefficients end-to-end. CorrSteer requires no contrastive data or additional storage, integrating token-level correlation analysis, mean-activation-driven coefficient estimation, and model steering. Evaluated on Gemma-2B and LLaMA-3.1-8B, it achieves a 4.1% absolute improvement in MMLU accuracy and a 22.9% gain in HarmBench safety using only 4,000 inference samples—demonstrating simultaneous enhancement of both task performance and safety across multiple benchmarks.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Safety and Robustness
• Search and Optimization: Metareasoning and Metaheuristics

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search
• Web Mining and Content Analysis: Large pretrained models with web data

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 avoiding 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 +4.1% improvement in MMLU performance and a +22.9% improvement in HarmBench with only 4000 samples. Selected features demonstrate semantically meaningful patterns aligned with each task's requirements, revealing the underlying capabilities that drive performance. Our work establishes correlationbased selection as an effective and scalable approach for automated SAE steering across language model applications.

Problem

• Improves LLM task performance via correlation-based feature selection
• Enhances safety by mitigating bias and preventing jailbreaking
• Automates steering using inference-time activations without contrastive datasets

Innovation

• Correlation-based feature selection for SAEs
• Uses inference-time activations only
• Automates steering coefficient extraction