GSAE: Graph-Regularized Sparse Autoencoders for Robust LLM Safety Steering

TL;DR

The paper tackles safety in large language models by reframing safety as a distributed, relational concept rather than a single latent axis. It introduces Graph-Regularized Sparse Autoencoders (GSAE), which incorporate a graph Laplacian on a neuron co-activation graph to learn smooth, distributed safety features and a spectral vector bank for steering. A dual-gated runtime controller selectively intervenes during generation, achieving high harmful-content refusal while preserving benign-task performance and generalizing across model families and jailbreak attacks. The findings demonstrate substantial safety gains with limited utility loss and establish a scalable, robust approach for real-time safety steering in diverse LLM deployments.

Abstract

Large language models (LLMs) face critical safety challenges, as they can be manipulated to generate harmful content through adversarial prompts and jailbreak attacks. Many defenses are typically either black-box guardrails that filter outputs, or internals-based methods that steer hidden activations by operationalizing safety as a single latent feature or dimension. While effective for simple concepts, this assumption is limiting, as recent evidence shows that abstract concepts such as refusal and temporality are distributed across multiple features rather than isolated in one. To address this limitation, we introduce Graph-Regularized Sparse Autoencoders (GSAEs), which extends SAEs with a Laplacian smoothness penalty on the neuron co-activation graph. Unlike standard SAEs that assign each concept to a single latent feature, GSAEs recover smooth, distributed safety representations as coherent patterns spanning multiple features. We empirically demonstrate that GSAE enables effective runtime safety steering, assembling features into a weighted set of safety-relevant directions and controlling them with a two-stage gating mechanism that activates interventions only when harmful prompts or continuations are detected during generation. This approach enforces refusals adaptively while preserving utility on benign queries. Across safety and QA benchmarks, GSAE steering achieves an average 82% selective refusal rate, substantially outperforming standard SAE steering (42%), while maintaining strong task accuracy (70% on TriviaQA, 65% on TruthfulQA, 74% on GSM8K). Robustness experiments further show generalization across LLaMA-3, Mistral, Qwen, and Phi families and resilience against jailbreak attacks (GCG, AutoDAN), consistently maintaining >= 90% refusal of harmful content.

Figures (6)

• Figure 1: Overview of the GSAE steering framework. A user query is encoded into hidden states, which the GSAE decomposes into graph-regularized safety features. A dual-gating controller uses these features to make a two-stage safety assessment: an Input Gate evaluates the initial prompt, while a Continuation Gate monitors the generation in real-time. This allows the system to selectively block harmful outputs while preserving benign ones.
• Figure 2: Safety performance across models, reported as the selective refusal score $\Delta_s$. GSAE (green) consistently outperforms both SafeSwitch (orange) and the baseline.
• Figure 3: Refusal trade-off plots: harmful refusal rate (HRR, y-axis) vs. safe refusal rate (SRR, x-axis). The ideal region is the top-left (maximizing harmful refusals while minimizing safe ones). GSAE consistently occupies this region, achieving the best balance.
• Figure 4: Distribution of safe vs. unsafe prompt activations projected onto the low-frequency eigenvectors of the neuron co-activation graph's Laplacian. The lack of a single eigenvector that cleanly separates the two distributions provides empirical support for the hypothesis that safety is a distributed concept.
• Figure 5: Distribution of per-feature Dirichlet energy for SAE vs. GSAE at an intermediate model layer. Both the PDF (left) and CDF (right) show that GSAE features (orange) are consistently smoother, possessing significantly lower energy than standard SAE features (blue). Dashed and dotted lines indicate the median energy for each model.