Interpretable Safety Alignment via SAE-Constructed Low-Rank Subspace Adaptation

TL;DR

This work tackles safety alignment for large language models by addressing polysemantic representations that impede identifying task-relevant subspaces. It introduces Sparse Autoencoders to disentangle activations into monosemantic features and constructs an explicit, SAE-derived subspace used to initialize the LoRA output projection, grounding updates in interpretable directions. Theoretical analysis shows SAE-based subspace recovery can achieve arbitrarily small error under monosemanticity, unlike direct polysemantic recovery, and empirically the method yields state-of-the-art safety performance with only ~0.2% of parameters updated, approaching RLHF baselines. The approach provides mechanistic interpretability for the learned alignment subspace and demonstrates robust improvements across model families and datasets, suggesting a practical path to safer, more transparent PEFT for LLMs.

Abstract

Parameter-efficient fine-tuning has become the dominant paradigm for adapting large language models to downstream tasks. Low-rank adaptation methods such as LoRA operate under the assumption that task-relevant weight updates reside in a low-rank subspace, yet this subspace is learned implicitly from data in a black-box manner, offering no interpretability or direct control. We hypothesize that this difficulty stems from polysemanticity--individual dimensions encoding multiple entangled concepts. To address this, we leverage pre-trained Sparse Autoencoders (SAEs) to identify task-relevant features in a disentangled feature space, then construct an explicit, interpretable low-rank subspace to guide adapter initialization. We provide theoretical analysis proving that under monosemanticity assumptions, SAE-based subspace identification achieves arbitrarily small recovery error, while direct identification in polysemantic space suffers an irreducible error floor. On safety alignment, our method achieves up to 99.6% safety rate--exceeding full fine-tuning by 7.4 percentage points and approaching RLHF-based methods--while updating only 0.19-0.24% of parameters. Crucially, our method provides interpretable insights into the learned alignment subspace through the semantic grounding of SAE features. Our work demonstrates that incorporating mechanistic interpretability into the fine-tuning process can simultaneously improve both performance and transparency.

Figures (3)

• Figure 1: Overview of our method. Top: SAEs transform polysemantic hidden states into a disentangled feature space where task-relevant directions become separable, enabling construction of an interpretable alignment subspace via PCA and QR decomposition. Bottom: Unlike LoRA which learns subspaces implicitly from data, we initialize $\mathbf{B}$ with semantically grounded basis vectors derived from SAE decoder directions.
• Figure 2: PCA visualization of SAE feature activations for aligned (blue) and unaligned (red) samples across different layers of Gemma 2 2B. Middle-deep layers (15--23) show the clearest separation, indicating concentrated encoding of safety-relevant concepts.
• Figure 3: Feature intervention experiments. Amplifying identified safety features reduces output toxicity, while suppressing them increases toxicity, demonstrating causal relevance.

Theorems & Definitions (13)

• Definition 1: Semantic Generative Model
• Definition 2: Task-Relevant Subspace
• Definition 3: Subspace Recovery Error
• Theorem 4: Original Space Recovery Error
• Theorem 5: SAE Space Recovery Error
• Theorem 6: Recovery Error Comparison
• Lemma 7: Original Space Differential
• proof
• proof : Proof of Theorem \ref{['thm:original_error']}
• Lemma 8: Feature Selection