AlignSAE: Concept-Aligned Sparse Autoencoders

TL;DR

This work addresses the opacity of knowledge in large language models by introducing AlignSAE, a concept-aligned sparse autoencoder trained in a two-stage curriculum on frozen LLM activations. By dedicating a small set of latent slots to ontology-defined relations and keeping the rest as a free feature bank, AlignSAE achieves near one-to-one binding in mid-layer representations and enables reliable causal interventions, such as concept swaps, with minimal interference to unrelated content. The approach yields a verifiable interface that supports slot-level inspection and steering, demonstrated on a synthetic biography QA task with strong generalization to unseen templates. Overall, AlignSAE provides a practical, modular step toward interpretable and controllable world-knowledge interfaces for frozen LLMs, laying groundwork for hierarchical ontologies, tool integration, and multi-hop reasoning in future work.

Abstract

Large Language Models (LLMs) encode factual knowledge within hidden parametric spaces that are difficult to inspect or control. While Sparse Autoencoders (SAEs) can decompose hidden activations into more fine-grained, interpretable features, they often struggle to reliably align these features with human-defined concepts, resulting in entangled and distributed feature representations. To address this, we introduce AlignSAE, a method that aligns SAE features with a defined ontology through a "pre-train, then post-train" curriculum. After an initial unsupervised training phase, we apply supervised post-training to bind specific concepts to dedicated latent slots while preserving the remaining capacity for general reconstruction. This separation creates an interpretable interface where specific relations can be inspected and controlled without interference from unrelated features. Empirical results demonstrate that AlignSAE enables precise causal interventions, such as reliable "concept swaps", by targeting single, semantically aligned slots.

Figures (7)

• Figure 1: An overview of our approach. Left: An unsupervised SAE trained post hoc on frozen LLM activations optimizes only reconstruction and sparsity, so each concept tends to be spread across multiple features, making interventions unreliable. Right: Our Concept-Aligned SAE adds a supervised binding loss that maps each concept to a dedicated feature, yielding clean, isolated activations that are easy to find, interpret, and steer.
• Figure 2: Comparison of binding generalization and causal intervention mechanisms.
• Figure 3: Relation–slot binding at a shallow layer (a) versus a mid layer (b) of GPT-2. At layer 0, supervision for each relation is dispersed across multiple slots, whereas at layer 6 the SAE learns a perfect one-to-one, diagonal binding, indicating that mid-layer representations are far more amenable to clean, controllable relation binding.
• Figure 4: Layer-wise concept fragmentation and concentration for AlignSAE (post-training) and a traditional SAE (pre-training only).
• Figure 5: Swap controllability across layers (rows) and amplification $\alpha$ (columns). Lighter is better; mid layers sustain robust control around $\alpha\!\approx\!2$.