CorrSteer: Il Pilotaggio Migliora le Prestazioni e la Sicurezza nei Modelli Linguistici di Grande Dimensione attraverso la Selezione di Caratteristiche basata su Correlazione con Autoencoder Sparse

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

August 18, 2025

Autori: Seonglae Cho, Zekun Wu, Adriano Koshiyama

cs.AI

Abstract

Gli Autoencoder Sparse (SAE) possono estrarre feature interpretabili da grandi modelli linguistici (LLM) senza supervisione. Tuttavia, la loro efficacia nei task di steering downstream è limitata dalla necessità di dataset contrastivi o di un ampio storage delle attivazioni. Per affrontare queste limitazioni, proponiamo CorrSteer, che seleziona le feature correlando la correttezza dei campioni con le attivazioni SAE dei token generati al momento dell'inferenza. Questo approccio utilizza solo le attivazioni al momento dell'inferenza per estrarre feature più rilevanti, evitando così correlazioni spurie. Inoltre, ottiene i coefficienti di steering dalle attivazioni medie, automatizzando l'intera pipeline. Il nostro metodo mostra un miglioramento delle prestazioni nei task di QA, mitigazione dei bias, prevenzione del jailbreaking e benchmark di ragionamento su Gemma 2 2B e LLaMA 3.1 8B, raggiungendo in particolare un miglioramento del +4,1% nelle prestazioni MMLU e del +22,9% in HarmBench con soli 4000 campioni. Le feature selezionate dimostrano pattern semanticamente significativi allineati con i requisiti di ciascun task, rivelando le capacità sottostanti che guidano le prestazioni. Il nostro lavoro stabilisce la selezione basata sulla correlazione come un approccio efficace e scalabile per lo steering automatizzato degli SAE nelle applicazioni dei modelli linguistici.

English

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.