AlignSAE: Concept-Aligned Sparse Autoencoders
summary
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
In short
AlignSAE upgrades Sparse Autoencoders (SAEs) to create interpretable interfaces for LLMs by enforcing alignment with a predefined ontology. It uses a two-phase training process—initial unsupervised learning followed by supervised post-training—to bind specific concepts to dedicated latent slots. This allows researchers to causally control model behavior through concept swaps, achieving significant disentanglement improvements.
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 used across episodes
This episode discusses
- AlignSAE: Concept-Aligned Sparse Autoencoders · Paper Radio
- SAEs Are Good for Steering -- If You Select the Right Features
- Steering Large Language Model Activations in Sparse Spaces
- Superposition as Lossy Compression: Measure with Sparse Autoencoders and Connect to Adversarial Vulnerability
- Towards Inference-time Category-wise Safety Steering for Large Language Models
- Improving Steering Vectors by Targeting Sparse Autoencoder Features
- Feature Hedging: Correlated Features Break Narrow Sparse Autoencoders
- On the Limits of Sparse Autoencoders: A Theoretical Framework and Reweighted Remedy
- Sparse Autoencoders Find Highly Interpretable Features in Language Models
- A Primer on the Inner Workings of Transformer-based Language Models
- Mechanistic Unlearning: Robust Knowledge Unlearning and Editing via Mechanistic Localization
- Llama Scope: Extracting Millions of Features from Llama-3.1-8B with Sparse Autoencoders
- LinguaLens: Towards Interpreting Linguistic Mechanisms of Large Language Models via Sparse Auto-Encoder
- Concept Bottleneck Models
- SAFER: Probing Safety in Reward Models with Sparse Autoencoder
- Decoupled Weight Decay Regularization
- Event Knowledge Incorporation with Posterior Regularization for Event-Centric Question Answering
- Sparse Feature Circuits: Discovering and Editing Interpretable Causal Graphs in Language Models
- Understanding Neural Networks via Feature Visualization: A survey
- Steering Language Model Refusal with Sparse Autoencoders
- Training language models to follow instructions with human feedback
The paper
AlignSAE: Concept-Aligned Sparse Autoencoders · Read on arXiv
University of Arizona · Peking University · Shandong University · LMU Munich
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.
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.
More episodes
- 2511.02406-Arithmetic Circuits and Neural Networks for Regular Matroids
- 2602.02583-CoCopula: Conformal Calibration and Copula Aggregation for Reliable Renewable Energy Forecasting
- 2601.21656-TabClustPFN: A Prior-data Fitted Network for Tabular Data Clustering
- 2310.17882-LOOP-MAC: Learning-Based Multi-Agent Coordination with Guaranteed Local Feasibility in a Virtual Power Plant
- 2405.00712-SoK: Behind the Accuracy of Complex Human Activity Recognition Using Deep Learning
- 2412.17907-A Multimodal Emotion Recognition System: Integrating Facial Expressions, Body Movement, Speech, and Spoken Language
- 2602.03737-Soft Sensor for Bottom-Hole Pressure Estimation in Petroleum Wells Using Long Short-Term Memory and Transfer Learning
- 2507.03411-Tourist arrival forecasting with game-theoretic social influence weighting and deep learning
- 2508.16531-Quality control in sublinear time: a case study via random graphs
- 2509.26226-Thinking-Free Policy Initialization Makes Distilled Reasoning Models More Effective and Efficient Reasoners