Technical Summary: Graph-Regularized Sparse Autoencoders for LLM Safety Steering

1. Problem Statement

Large Language Models (LLMs) possess strong capabilities but remain vulnerable to adversarial prompts and jailbreaks that elicit harmful generations. Effective safety mechanisms must balance two competing requirements: reliably blocking harmful outputs while preserving utility on benign and dual-use queries.

Current defenses face significant limitations:

• Black-box guardrails (e.g., prompt filtering, output classifiers) operate outside the model's internal computation, making them brittle under distribution shifts and lacking adaptive robustness.
• Internals-based methods (e.g., activation additions, refusal directions) intervene directly on hidden representations. However, many of these approaches, including standard Sparse Autoencoders (SAEs), treat latent features as independent. This independence prior is poorly matched to high-level safety behaviors like refusal and harmful compliance, which recent studies suggest depend on distributed, nonlinear structures in activation space rather than isolated semantic axes.

The core problem addressed is that standard SAEs may fragment coherent safety concepts across weakly aligned features, reducing the stability and effectiveness of steering directions derived from them.

2. Methodology

The authors propose Graph-Regularized Sparse Autoencoders (GSAE), a framework that learns smoother, structurally coherent safety-steering directions by incorporating the topology of neuron co-activations. The method consists of three phases:

2.1 Graph-Regularized Training (GSAE)

The authors extend the standard SAE objective by adding a Laplacian smoothness penalty over a neuron co-activation graph.

• Graph Construction: An undirected graph $G=(V, E)$ is constructed where nodes represent neurons in a target layer. Edge weights are defined by the thresholded cosine similarity of neuron activation profiles across a held-out set of safety-labeled prompts.
• Objective Function: The training loss combines reconstruction error, sparsity ($L_1$), supervised classification loss (via a linear probe), and the new graph regularization term:
  $$L_{GSAE} = \|h - \hat{h}\|^2_2 + \lambda_{spar}\|z\|_1 + \lambda_{sup}L_{cls} + \lambda_{graph} \sum_{j=1}^k (W^{(d)}_{\cdot, j})^\top L W^{(d)}_{\cdot, j}$$
  Here, $L$ is the symmetric normalized Laplacian of the co-activation graph. The graph term penalizes decoder directions that vary sharply across strongly co-activating neurons, biasing the learned dictionary toward low-frequency eigenmodes that respect the distributed structure of safety-relevant patterns.

2.2 Steering Bank Construction

Not all learned features are suitable for steering. The authors construct a curated "steering bank" by scoring decoder directions against three independent criteria:

1. Structural Coherence ($s_{coh}$): Measures the smoothness of a direction on the neuron graph (low Dirichlet energy).
2. Predictive Relevance ($s_{rel}$): The magnitude of the coefficient in the linear probe trained to predict safety labels.
3. Intervention Efficacy ($s_{eff}$): The actual change in model behavior (harm reduction vs. over-refusal) when the direction is applied.

Directions are scored using a weighted geometric mean of these metrics. Only directions that perform well on all three are retained in the final bank.

2.3 Runtime Steering with Two-Gate Controller

To prevent degradation of benign performance, GSAE employs a selective steering pipeline with two gates:

• Input Gate: Evaluates the prompt-level risk score before generation. If risk is high, it triggers an immediate refusal; if low, generation proceeds without steering.
• Continuation Gate: Monitors token-level risk during generation. It activates steering only when risk rises above a high threshold and deactivates it only when risk falls below a low threshold (using hysteresis to prevent unstable switching).
• Intervention: When active, the controller applies a corrective shift to hidden states by subtracting a weighted sum of the steering bank directions, scaled by their alignment with the current hidden state.

3. Key Contributions

1. GSAE Objective: Introduction of a graph-regularized SAE objective that enforces smoothness over a neuron co-activation graph, addressing the limitation of independence priors in standard SAEs.
2. Selective Steering Pipeline: Development of a multi-criteria bank construction method (coherence, relevance, efficacy) combined with a two-gate runtime controller (input and continuation) to ensure selective intervention.
3. Empirical Isolation: Controlled ablation studies demonstrating that graph regularization specifically improves the quality of SAE-derived intervention directions, distinct from the benefits of the downstream gating or bank construction.

4. Experimental Results

The method was evaluated primarily on Llama-3-8B-base, with cross-model validation on Mistral 7B, Qwen 2.5 14B, and Phi-4 15B.

• Safety Performance: On JailbreakBench, replacing a standard SAE with GSAE (in an otherwise identical pipeline) improved the selective refusal score ($\Delta s = HRR - SRR$) by 20.1 points (from 58.8% to 78.9%). Similar gains were observed on HarmBench (+16.8 points) and XSTest.
• Comparison to Baselines: GSAE outperformed activation-steering baselines (e.g., CAA, SafeSwitch) and black-box guardrails (e.g., Llama Guard input+output) while maintaining comparable or better throughput.
• Utility Preservation: The method preserved benign-task performance, with only a modest average drop (4.2 points) in utility benchmarks (TriviaQA, GSM8K, TruthfulQA) compared to the unsteered baseline.
• Generalization: The performance gains transferred across different model architectures (Mistral, Qwen, Phi).
• Robustness: GSAE remained strong under black-box (AutoDAN, TAP) and gray-box (GCG) jailbreak attacks. However, under a steering-aware white-box stress test (where the attacker has access to the bank vectors and gate thresholds), robustness degraded, with Harmful Refusal Rate (HRR) dropping to 51.9%.
• Mechanism Analysis: Diagnostic probes indicated that GSAE activates explicit refusal trajectories rather than merely shifting hidden states toward a benign-looking region (safe-centroid shift).

5. Significance and Claims

The paper claims that graph regularization is a critical inductive bias for learning safety-steering directions from SAEs. The primary finding is that safety behaviors in LLMs are distributed and nonlinear; enforcing structural smoothness over co-activating neurons prevents the fragmentation of safety concepts, leading to more effective intervention directions.

The authors emphasize that their contribution is controlled and functional: they isolate the specific improvement gained from the graph term under a matched downstream pipeline. They do not claim to have discovered the complete mechanistic circuitry of refusal, but rather that graph-regularized directions provide a superior basis for selective steering.

The work highlights a practical deployment trade-off: the full two-gate controller offers token-level precision at a modest latency cost, while an input-gate-only variant retains much of the safety gain with lower overhead. The authors acknowledge limitations, noting that the method requires hidden-state access (open-weight models), relies on a fixed co-activation graph that may not adapt to distribution shifts, and is vulnerable to white-box attacks that expose the steering internals.