AlignSAE: Concept-Aligned Sparse Autoencoders

Minglai Yang; Xinyu Guo; Zhengliang Shi; Jinhe Bi; Steven Bethard; Mihai Surdeanu; Liangming Pan

cs.LGDec 1, 2025

AlignSAE: Concept-Aligned Sparse Autoencoders

Authors: Minglai Yang, Xinyu Guo, Zhengliang Shi, Jinhe Bi, Steven Bethard, Mihai Surdeanu, Liangming Pan

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

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$ -- $\infty\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.

Figures & tables

Figure 1: Overview of AlignSAE . 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 AlignSAE adds a supervised binding loss that maps each concept in a pre-defined ontology to a dedicated feature, yielding clean, isolated activations that are easy to find, interpret, and control.

Figure 2: 2-hop step-wise alignment. With

(e_{1}\mathord{\mathchar 59\relax}r_{1}\mathord{\mathchar 59\relax}r_{2})

, the model predicts

(e_{2}\mathord{\mathchar 59\relax}e_{3})

;

r_{1}

activates at

e_{2}

and

r_{2}

e_{3}

, with remaining capacity in unsupervised features.

Figure 3: Concept–slot binding at a shallow layer (a) versus a mid layer (b). At layer 0, supervision for each relation is dispersed across multiple slots, whereas at layer 6 the SAE learns a perfect one-to-one binding.

Figure 4: Binding generalization vs. causal control across layers. Left: diagonal accuracy on train (blue) and on unseen templates (orange). Right: diagonal accuracy and maximum swap success over

\alpha

, both on held-out test questions.

Metric	Layer 0	Layer 6	$\Delta$ (L6 – L0)
Diagonal Acc (test) $\uparrow$	0.238	1.000	$\uparrow\ 0.76$
Swap Success $\uparrow$	0.040	0.847	$\uparrow\ 0.81$
Diagonal Acc (train) $\uparrow$	0.232	1.000	$\uparrow\ 0.77$
Diagonal Acc (unseen) $\uparrow$	0.165	0.912	$\uparrow\ 0.75$
Recon MSE $\downarrow$	$6.53\cdot 10^{-5}$	$7.42\cdot 10^{-2}$	$\uparrow\ \approx 1.1\text{k}\times$

Table 1: Layer 0 vs. Layer 6 performance comparison.

Panel of Figure 5 — Figure 5: Layer-wise fragmentation ( $\downarrow$ ) and concentration ( $\uparrow$ ) for AlignSAE and a traditional SAE.

Figure 6: Swap controllability across layers and amplification $\alpha$ . Mid layers remain robust at $\alpha\!\approx\!2$ .

Target swap	$\alpha{=}2$ (15% Err)			$\alpha{=}10$ (77% Err)
Target swap	S	D	S%	S	D	S%
birth_city	36	20	64.3	91	50	64.5
birth_date	20	0	100.0	139	0	100.0
employer	8	1	88.9	134	3	97.8
major	1	0	100.0	77	51	60.2
university	42	10	80.8	101	1	99.0

Table 2: Category retention on failed swaps at Layer 6. “Same” (S) means correct semantic class; “Diff” (D) is outside.

Swap	Q: Where did Jesse Kian Tate go to college?
	univ $\rightarrow$ Swap: major
Outputs	Base: Rochester Inst. of Tech.
	Gold: Physical Therapy
	Gen: Geography [ Type Entity ]

Table 3: A failure case: steering (

\alpha=10

) successfully flips the semantic category (type) but fails to retrieve the specific ground-truth entity.

Config	Train Diag. Acc	Test-Unseen Diag. Acc	Train Acc	Test-Unseen Acc	Swap Success Rate
joint	1.000	0.912	1.000	0.809	0.847
no_align	0.167	0.167	1.000	0.809	0.053
no_ortho	0.711	0.689	1.000	0.809	0.046
no_value	0.177	0.172	1.000	0.809	0.720

Table 4: Ablations of AlignSAE post-training. Best configuration is joint , which achieves strong OOD binding and high swap success simultaneously. Swap success is reported at each configuration’s best-performing intervention strength (

\alpha

: joint=2, no_align=0.1, no_ortho=10, no_value=100). Chance diagonal accuracy for 6 relations is

1/6\approx 0.167

Metric	No Stage 1	With Stage 1	$\Delta$ (With–No)	% $\Delta$
$\mathcal{L}\downarrow$	3.3438	1.8594	$L\downarrow~1.4844$	-44.4%
$L_{\text{recon}}\downarrow$	0.0412	0.0256	$L\downarrow~0.0156$	-37.9%
$L_{\text{sparse}}\downarrow$	0.2682	0.0209	$L\downarrow~0.2473$	-92.2%
$L_{\text{align}}\downarrow$	0.0605	0.0762	$L\uparrow~0.0157$	+26.0%
$L_{\text{ortho}}\downarrow$	0.5330	0.3760	$L\downarrow~0.1570$	-29.5%
$L_{\text{value}}\downarrow$	3.3169	3.4125	$L\uparrow~0.0956$	+2.9%

Table 5: Stage-1 ablation. Deltas are computed as With–No ; losses prefer

\downarrow

and accuracies prefer

\uparrow

Figure 7: Swap controllability in 2-hop reasoning. The post-trained AlignSAE achieves

4\times

higher swap success than the traditional SAE when

\alpha\approx 50

Figure 8: Relation-to-slot binding (pre- vs. post-grokking). At steps $e_{2}$ and $e_{3}$ , binding is initially entangled but becomes perfectly diagonal after grokking, indicating clean step-wise alignment.

Figure 9: Grokking-like emergence. Diagonal accuracy saturates early, while validation accuracy improves later, suggesting delayed generalization in 2-hop reasoning.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Relation	Train templates (T0–T1)	Held-out templates (T2–T3)
Birth Date	When was {FULL_NAME} born? On what date was {FULL_NAME} born?	What is {FULL_NAME}’s birth date? Can you tell me the birth date of {FULL_NAME}?
Birth City	Where was {FULL_NAME} born? In what city was {FULL_NAME} born?	What is {FULL_NAME}’s birth city? Can you tell me the birth city of {FULL_NAME}?
University	Where did {FULL_NAME} go to college? Which college did {FULL_NAME} attend?	What is {FULL_NAME}’s alma mater? Which university did {FULL_NAME} attend?
Major	What was {FULL_NAME}’s major? What is {FULL_NAME}’s field of study?	What did {FULL_NAME} study? What field did {FULL_NAME} study in?
Employer	Who does {FULL_NAME} work for? What company does {FULL_NAME} work for?	What is {FULL_NAME}’s employer? Which company employs {FULL_NAME}?
Work City	Where does {FULL_NAME} work? What city does {FULL_NAME} work in?	What is {FULL_NAME}’s work city? In which city is {FULL_NAME} employed?

Appendix

Table 6: Template split for surface-form robustness. T0–T1 are used for training, while T2–T3 are held out for Unseen-Template evaluation; the split is designed to reduce lexical overlap in relation triggers ( e.g. , born vs. birth date , work for vs. employer ).

Panel of Figure 11 — Figure 11: Layer-wise top-50 feature activations in GPT-2 with and without SAE post-training. Early layers (0–3) remain diffuse and entangled. Middle layers (4–8) show maximum SAE benefit, with Layer 6 achieving optimal semantic binding and diagonal accuracy. Deep layers (9–11) without SAEs exhibit irregular patterns and over-compression, whereas SAE-trained models maintain structured representations.

Figure 12: Illustrative schematic of our swap intervention (Orig relation $\rightarrow$ Swap relation).

Orig $\rightarrow$ Swap	Question	Target	Generated
company_city $\rightarrow$ university	What is Grace Wendy Rivera’s work city?	Florida International University	Florida International University
company_city $\rightarrow$ major	What is Grace Wendy Rivera’s work city?	Electrical Engineering	Electrical Engineering
company_city $\rightarrow$ employer	What is Grace Wendy Rivera’s work city?	Blackstone	Blackstone
university $\rightarrow$ birth_date	Where did Thomas Heath Stafford go to college?	2, March, 1981	2, March, 1981
university $\rightarrow$ major	Where did Thomas Heath Stafford go to college?	Dance	Dance
birth_date $\rightarrow$ work_city	When was Megan Kian Valencia born?	Framingham, MA	Framingham, MA

Appendix

Table 7: Correct swap examples from the evaluation set (Layer 6;

\alpha{=}2

Setting	AlignSAE	Standard SAE	Ratio
1-Hop (6 concepts)	0.173	0.000	$\infty$
2-Hop (20 relations)	0.035	0.004	$9.2\times$

Appendix

Table 8: Mutual Information Gap (MIG). Standard SAE uses identical architecture trained with only reconstruction and sparsity losses. An MIG of 0.000 is expected for the standard SAE: without concept supervision, no single unsupervised feature has any incentive to exclusively encode a specific concept—information is distributed across many features by design.

Panel of Figure 13 — Figure 13: t-SNE projections of raw activations (left) vs. AlignSAE supervised slot activations (right), colored by concept label. Top: 1-hop (6 concepts). Bottom: 2-hop (20 relations).

Figure 14: Grokking progression: t-SNE of 2-hop supervised slot activations at epochs 5, 10, 25, 50, 100.

Figure 15: Perplexity vs. steering strength

\alpha

. Correct-slot steering (blue) dramatically reduces perplexity; wrong-slot steering (red) does not. Dashed lines show no-SAE and SAE-reconstruct baselines.

Condition	First tokens generated
No steering	Whitman College counts She among its graduates. She works in Purchase, NY. Her field of study was Nursing…
Steer $\to$ birth_date ( $\alpha{=}10$ )	27,July,1991 , Bryan Caleb Atkins was born. He is professionally active in Fort Worth, TX…
Steer $\to$ major ( $\alpha{=}10$ )	French . According to records, Maurice Hudson Tran was born on 28,September,1953…
Steer $\to$ work_city ( $\alpha{=}10$ )	Dallas, TX . He earned a degree in Nursing. He is a key part of the AT&T team…

Appendix

Table 9: Free-form generation examples (1-hop,

\alpha{=}10

). Bold text marks the first tokens generated, which immediately reflect the steered concept. Continuations remain fluent biographical text.

Setting	Condition	Fluency	Coherence	Steering	$n$
1-Hop	Baseline	2.0	1.0	—	5
	Steered $\alpha{=}10$	2.1	1.1	3.9	15
	Steered $\alpha{=}100$	2.0	1.1	3.9	15
2-Hop	Baseline	4.0	1.2	—	5
	Steered $\alpha{=}10$	4.0	1.3	2.5	15
	Steered $\alpha{=}100$	3.7	1.3	2.4	15

Appendix

Table 10: LLM-as-Judge evaluation of free-form generation quality (1–5 scale). Steering preserves fluency and coherence while successfully injecting the target concept in 1-hop.

Metric	Value
Train slot top-1 accuracy	0.987
Test slot top-1 accuracy (OOD)	0.757
Reconstruction MSE	0.083
Diagonal accuracy	0.956

Appendix

Table 11: AlignSAE performance in the 91-relation FB15K setting.

Figure 16: Normalized confusion matrix for 91-relation slot binding at Layer 6. The strong diagonal confirms that AlignSAE maintains near-perfect concept–slot alignment even at this scale.

Idx	Relation	#Ent	Idx	Relation	#Ent
0	/award/award_category/category_of	90	45	/location/country/second_level_divisions	11
1	/award/award_category/disciplines_or_subjects	90	46	/location/hud_county_place/county	90
2	/base/aareas/schema/administrative_area/ administrative_area_type	90	47	/location/hud_county_place/place	90
3	/base/aareas/schema/administrative_area/ administrative_parent	90	48	/location/location/contains	90
4	/base/americancomedy/celebrity_impressionist/ celebrities_impersonated	8	49	/location/location/partially_contains	90
5	/base/biblioness/bibs_location/country	90	50	/location/location/time_zones	90

Appendix

Table 12: All 91 FB15K relations with the number of unique head entities per relation. Relations with fewer than 90 entities retain their original count.

Figure 18: Training dynamics on the two-hop task (base model). Top: training vs. validation loss. Bottom: validation accuracy for the 1st hop (Token 1, predicting the intermediate entity $e_{2}$ ) and the 2nd hop (Token 2, predicting the final entity $e_{3}$ ) under step-wise supervision.

Sparse Autoencoders

Sparse autoencoders decompose neural network activations into sparse features to study, interpret, and control learned representations.

See all 218 papers

Navigating the Concept Space of Language Models
Mar 6, 2026Wilson E. Marcílio-Jr, Danilo M. ElerLLM Interpretability Representation Geometry in Language Models
Adaption Labs, Curitiba, PR, Brazil · São Paulo State University (UNESP), Presidente Prudente, SP, Brazil
Step-Level Sparse Autoencoder for Reasoning Process Interpretation
Mar 3, 2026Xuan Yang, Jiayu Liu, Yuhang Lai +3LLM Interpretability Sparse Autoencoders
Department of Data Science, City University of Hong Kong · Hong Kong Institute of AI for Science, City University of Hong Kong · Li Auto Inc +1
SAEVerbalizer: Generating Explanations for Sparse Autoencoder Features via Representation Verbalization
Aug 13, 2026Weihan Meng, Hongzhu Guo, Yi Jing +5LLM Interpretability Sparse Autoencoders
Tsinghua University · Peking University · Fudan University

Mar 6, 2026J/K move · Enter open · S save

Navigating the Concept Space of Language Models

Wilson E. Marcílio-Jr, Danilo M. Eler

Adaption Labs, Curitiba, PR, Brazil · São Paulo State University (UNESP), Presidente Prudente, SP, Brazil

LLM Interpretability Representation Geometry in Language Models

Open paper AskPDF arXiv

Sparse autoencoders (SAEs) trained on large language model activations output thousands of features that enable mapping to human-interpretable concepts. The current practice for analyzing these features primarily relies on inspecting top-activating examples, manually browsing individual features, or performing semantic search on interested concepts, which makes exploratory discovery of concepts difficult at scale. In this paper, we present Concept Explorer, a scalable interactive system for post-hoc exploration of SAE features that organizes concept explanations using hierarchical neighborhood embeddings. Our approach constructs a multi-resolution manifold over SAE feature embeddings and enables progressive navigation from coarse concept clusters to fine-grained neighborhoods, supporting discovery, comparison, and relationship analysis among concepts. We demonstrate the utility of Concept Explorer on SAE features extracted from SmolLM2, where it reveals coherent high-level structure, meaningful subclusters, and distinctive rare concepts that are hard to identify with existing workflows.

Abstract

Sparse Autoencoders

Explore similar work

Navigating the Concept Space of Language Models

Step-Level Sparse Autoencoder for Reasoning Process Interpretation

SAEVerbalizer: Generating Explanations for Sparse Autoencoder Features via Representation Verbalization