CorrSteer: La Dirección Mejora el Rendimiento y la Seguridad en LLMs mediante la Selección de Características basada en Autocodificadores Dispersos y Correlación

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

August 18, 2025

Autores: Seonglae Cho, Zekun Wu, Adriano Koshiyama

cs.AI

Resumen

Los Autoencoders Escasos (SAEs, por sus siglas en inglés) pueden extraer características interpretables de modelos de lenguaje grandes (LLMs) sin supervisión. Sin embargo, su eficacia en tareas de control posteriores se ve limitada por la necesidad de conjuntos de datos contrastivos o un gran almacenamiento de activaciones. Para abordar estas limitaciones, proponemos CorrSteer, que selecciona características correlacionando la corrección de las muestras con las activaciones de los SAEs a partir de tokens generados en tiempo de inferencia. Este enfoque utiliza únicamente activaciones en tiempo de inferencia para extraer características más relevantes, evitando así correlaciones espurias. Además, obtiene coeficientes de control a partir de activaciones promedio, automatizando todo el proceso. Nuestro método muestra un mejor rendimiento en tareas de preguntas y respuestas (QA), mitigación de sesgos, prevención de jailbreaking y pruebas de razonamiento en Gemma 2 2B y LLaMA 3.1 8B, logrando notablemente una mejora del +4.1% en el rendimiento de MMLU y un +22.9% en HarmBench con solo 4000 muestras. Las características seleccionadas demuestran patrones semánticamente significativos alineados con los requisitos de cada tarea, revelando las capacidades subyacentes que impulsan el rendimiento. Nuestro trabajo establece la selección basada en correlación como un enfoque efectivo y escalable para el control automatizado de SAEs en aplicaciones de modelos de lenguaje.

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.