AlignSAE: Concept-Aligned Sparse Autoencoders

arXiv:2512.02004 · cs.LG, cs.CL · Submitted 2025-12-01 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "AlignSAE: Concept-Aligned Sparse Autoencoders".

Jane: The gist The ALIGNSAE framework upgrades Sparse Autoencoders from descriptive probes to operational interfaces by introducing a “pre-train,

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So we're looking at the paper titled AlignSAE: Concept-Aligned Sparse Autoencoders. It sounds like they are trying to fix a problem where these sparse autoencoders, which are supposed to break down what an AI is thinking into smaller pieces, end up being messy and mixed up.

Jane: Exactly. The title tells us the goal is to get these autoencoders aligned with some kind of defined knowledge structure, an ontology. It’s about taking those hidden layers of the large language model and making them interpretable by linking specific features to specific concepts we can actually name, like "birth date" or "city."

Lu: The core idea here is that standard SAE training just optimizes for reconstruction and sparsity in an unsupervised way, so each concept gets spread out across many different features. This makes it hard to know which feature controls which concept.

Meng: So, what's the mechanism they propose to fix this entanglement? They introduce a two-phase training curriculum called pre-train then post-train. That sounds like they train it in one way first, and then fine-tune it with a specific goal later.

Lalam: The paper introduces this "pre-train, then post-train" curriculum to enforce that explicit alignment between the designated latents and an external ontology. It's not just about reconstruction anymore; they add supervised post-training to bind specific concepts to dedicated latent slots while keeping enough capacity left over for general reconstruction.

Tom: So they’re essentially doing an initial unsupervised training phase, and then applying supervised post-training to lock down those specific concepts into their own features. That separation is key because it creates an interface where you can inspect and control a concept without it getting mixed up with unrelated features.

The paper's summary: Jane: So, if we boil down the summary of AlignSAE, they are tackling the fact that standard SAEs don't naturally map their features to human-defined concepts. They show that by using this pre-train, then post-train method, you can create sparse representations where each concept has its own dedicated slot.

Lu: The paper points out that in the unsupervised training phase, a target concept might be fragmented across multiple features, and individual features might mix unrelated signals. That’s what they are trying to avoid when we talk about reliable control.

Tom: Right, so the improvement isn't just making it sparse; it's adding this supervised binding loss to map each concept in a pre-set ontology to a dedicated feature. This means we get clean, isolated activations that are easy to find, interpret, and steer.

Meng: So what does this mean for deployment? If you have an AI that’s using these SAEs, instead of looking at hundreds of features for one concept like "major work city," you can look directly at the feature dedicated to that city. That sounds much more manageable for debugging issues at scale.

Lalam: The paper suggests this separation creates a Verifiable Interface where you can check whether the model is using a particular relation by observing whether the corresponding concept slot activates, and then you can causally steer the computation by manually activating or suppressing that slot.

Jane: So, it moves SAEs from being descriptive probes to actual operational interfaces. It gives us reliable causal control over what the AI is doing because we know which feature corresponds to which piece of knowledge.

The paper's improvements: Tom: The paper outlines several improvements they introduce, focusing on how this alignment actually translates into usable control. They are looking at things like concept–slot binding, layer-wise performance, and swap controllability.

Lu: They show that the concept–slot binding becomes clean and diagonal at mid layers when using AlignSAE, whereas a traditional SAE keeps features highly fragmented across all layers. This means the representations are much more compact from mid layers onward.

Meng: That fragmentation issue is a big problem for production because it suggests we can't trust the features reliably deep in the model structure. If you look at their layer-wise performance, they show that AlignSAE sharply reduces fragmentation and increases concentration via post-training supervision.

Jane: And they test if these slots actually work by measuring swap controllability. They find that targeting single, semantically aligned slots allows for precise causal interventions, achieving an optimal swap success rate of eighty-five percent at layer six.

Lalam: The results show that this concept–slot binding is robust and generalizes well to unseen templates. They also use this framework to probe grokking dynamics in compositional tasks, tracking how concept evidence concentrates into structured, compositional features as the model moves from an entangled regime to a clean, compositional regime.

Conclusion: Tom: So wrapping up the AlignSAE paper, they're showing that by adding this supervised binding loss and using their pre-train then post-train approach, you get representations that are not only sparse but also explicitly aligned with concepts.

Jane: They demonstrate that this results in better interpretability and disentanglement than a standard SAE, especially when looking at layer-wise performance and how well the concepts bind to features. It really shows how a small change in the training objective can have such a big impact on usability.

Lu: The implication for representation learning is that we can get representations that are much more structured, moving from an entangled state to one where concepts are localized to single features, which supports reliable steering and control.

Meng: From an engineering standpoint, this means we can deploy systems where we aren't guessing which feature is active; we can actually intervene on specific pieces of knowledge with high certainty. That level of control is something you need when you’re trying to build something that needs to be dependable performance.

Lalam: The AlignSAE framework allows for reliable concept swaps by targeting those dedicated slots, and it also provides a direct, verifiable readout of the concept in use through observing which slot activates, which is really powerful for mechanistic interpretability.

Tom: So we have this paper, AlignSAE: Concept-Aligned Sparse Autoencoders. It takes sparse autoencoders and makes them operational interfaces by enforcing alignment with an ontology through that pre-train then post-train curriculum. It’s about getting concrete control over the AI's internal workings so we can actually use them reliably.

Jane: It’s a solid direction for making sure the features we extract from LLMs are meaningful and not just noise. We'll be looking at how this idea plays out in more complex reasoning tasks next time.

University of Arizona · Peking University · Shandong University · LMU Munich

cs.LG, cs.CL

Submitted: 2025-12-01

Updated: 2026-10-15

Comments: Published in Transactions on Machine Learning Research (08/2026). 36 pages, 18 figures, 12 tables. v2: adds an Errata section clarifying a few issues in the TMLR version

Journal ref: Transactions on Machine Learning Research (TMLR), 2026. https://openreview.net/forum?id=I9UjKxW4nq

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 88/100

The gist: The gist The ALIGNSAE framework upgrades Sparse Autoencoders from descriptive probes to operational interfaces by introducing a “pre-train, then post-train” curriculum that enforces explicit

Key concepts

Sparse Autoencoders (SAEs)
SAEs are a technique used to decompose the complex hidden activations inside Large Language Models into a smaller set of meaningful, sparse features. The goal is to find a compact representation where each feature relates clearly to some aspect of the model's internal knowledge.
Ontology Alignment
This refers to mapping specific, human-defined concepts (like 'relation A' or 'concept B') onto dedicated, isolated latent slots within the SAE. Instead of features being entangled randomly, this process forces specific concepts to occupy unique spaces that can be inspected and controlled independently.
Pre-train, then Post-train Curriculum
This is a two-step training method. First, the SAE is trained unsupervised to learn general reconstruction and sparsity. Second, a supervised phase is added where the model learns to explicitly bind specific relations to their designated concept slots while ensuring other features remain free for general reconstruction.

Terminology

Summary

The gist The ALIGNSAE framework upgrades Sparse Autoencoders from descriptive probes to operational interfaces by introducing a “pre-train, then post-train” curriculum that enforces explicit alignment between designated latents and an external ontology, yielding reliable causal control and interpretable representations.

Introduction

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 finegrained, interpretable features, they often struggle to reliably align these features with humandefined concepts, resulting in entangled and distributed feature representations. To address this, we introduced ALIGNSAE, a method that aligns SAE features with a predefined 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 concepts 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, and further supports multi-hop reasoning and a mechanistic probe of grokking-like generalization dynamics.

Concept Binding in Activation Space

Given a frozen LLM M, we extract an activation h from layer l and learn a sparse representation z = E(h) ∈ R R+K with an SAE z = [zconcept; zrest], where zconcept are R concept slots and zrest are K free features (R ≪ K). This makes concept–slot activations a direct, verifiable readout of the concept in use. A Verifiable Interface is constructed where (i) verification checks whether the model is using a particular relation by observing whether the corresponding concept slot activates; and (ii) control allows us to causally steer the computation by manually activating or suppressing that slot, directly influencing the model’s downstream prediction.

Model Training

The training protocol involves a two-phase curriculum that parallels LLM pre-training: first unsupervised, then supervised. In the pre-training phase, the SAE is trained on reconstruction and sparsity. In subsequent post-training, we strengthen the binding and value losses, then activate the orthogonality penalty. The total training objective (see Section 4 of the main paper) combines six components. These objectives are LSAE = λrec∥h − hˆ∥2 + λsp∥z∥1, Lalign = CEsoftmax(zconcept), yrel, and L⊥ = corr(zconcept, zrest). The binding objective augments the standard SAE loss (LSAE) with (i) a supervised loss that assigns each relation to a dedicated concept slot (Lalign), (ii) a zconcept–zrest decorrelation penalty to reduce concept leakage into zrest (L⊥), and (iii) an auxiliary value loss that encourages the concept slots to support answer prediction.

Empirical Results on Factual Recall

Empirically, ALIGNSAE is more interpretable, disentangled, and controllable than a standard SAE under three RC evaluations. (i) Binding/generalization: concept–slot binding becomes clean and diagonal at mid layers (Figure 4) and transfers to unseen templates (Figure 3a).

Layer-Wise Performance and Concept–Feature Alignment

Figure 5 reports layer-wise fragmentation (↓) and concentration (↑) for ALIGNSAE and a traditional SAE. The traditional SAE remains highly fragmented across all layers: EffFeat spans hundreds to thousands of features per concept and Top1C ≈ 0, suggesting that concepts are not consistently localized to a single feature. In contrast, ALIGNSAE sharply reduces fragmentation (EffFeat ≈ 1) and increases concentration via post-training supervision, yielding compact representations from mid layers onward and near oneto-one bindings in deeper layers.

Swap Controllability

To test whether ontology-aligned slots serve as usable control knobs, we measure the success rate of answer swaps under inference-time steering along decoded SAE directions. The swap intervention accuracy validates that these representations causally influence the full model’s behavior. Optimal swap success: 85% at α ≈ 2 (Layer 6).

Grokking and Concept Binding

We use ALIGNSAE as a probe to understand grokking in compositional tasks.

Improvements for AI systems

  1. A mechanism for precise causal interventions via concept swaps can be implemented by targeting single, semantically aligned slots to reliably control model outputs, as demonstrated by achieving a 4× higher swap success than the traditional SAE (Figure 8).

  2. The system can perform reliable multi-hop reasoning by utilizing the mechanism where posttrained ALIGNSAE attains 4× higher swap success than the traditional SAE (Figure 8) and achieves 100.0% Binding Accuracy at both step decoding tokens in two-hop tasks (Figure 7).

  3. The AI can exhibit mechanistic interpretability by providing a direct, verifiable readout of the concept in use through observing which slot is active, as the paper demonstrates that we can check whether the model is using a particular relation by observing whether the corresponding concept slot activates.

  4. The system can be used to diagnose grokking dynamics in compositional tasks by tracking how concept evidence concentrates into structured, compositional features that support stable step-wise 1-to-1 concept binding as the model moves from an entangled regime to a clean, compositional regime (Figure 16).

  5. The system can be improved for robustness against surface form variations in factual recall tasks by using the template split strategy, which is designed to reduce n-gram overlap with training prompts to assess robustness to surface form (Table 4).

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 predefined 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 concepts can be inspected and controlled without interference from unrelated features. Empirical results on GPT-2 across synthetic biographical QA (1-hop, 6 concepts) and compositional relational reasoning (2-hop, 20 concepts) demonstrate that AlignSAE enables precise causal interventions, such as reliable concept swaps (94.8% success in 1-hop), by targeting single, semantically aligned slots. AlignSAE achieves disentanglement improvements of 9 -- infinity times over standard SAEs (measured via Mutual Information Gap), and further supports multi-hop reasoning and a mechanistic probe of grokking-like generalization dynamics. Although our evaluation focuses on GPT-2, AlignSAE itself is model-agnostic.

Sources

Related papers