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--∞× 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 (e1\mathchar59r1\mathchar59r2) , the model predicts (e2\mathchar59e3) ; r1 activates at e2 and r2 at e3 , 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 α , both on held-out test questions.
Metric
Layer 0
Layer 6
Δ (L6 – L0)
Diagonal Acc (test) ↑
0.238
1.000
↑0.76
Swap Success ↑
0.040
0.847
↑0.81
Diagonal Acc (train) ↑
0.232
1.000
↑0.77
Diagonal Acc (unseen) ↑
0.165
0.912
↑0.75
Recon MSE ↓
6.53⋅10−5
7.42⋅10−2
↑≈1.1k×
Table 1: Layer 0 vs. Layer 6 performance comparison.
Figure 5: Layer-wise fragmentation ( ↓ ) and concentration ( ↑ ) for AlignSAE and a traditional SAE.
Figure 6: Swap controllability across layers and amplification α . Mid layers remain robust at α≈2 .
Target swap
α=2 (15% Err)
α=10 (77% Err)
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 → Swap: major
Outputs
Base: Rochester Inst. of Tech.
Gold: Physical Therapy
Gen: Geography [ Type Entity ]
Table 3: A failure case: steering ( α=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 ( α : joint=2, no_align=0.1, no_ortho=10, no_value=100). Chance diagonal accuracy for 6 relations is 1/6≈0.167 .
Metric
No Stage 1
With Stage 1
Δ (With–No)
% Δ
L↓
3.3438
1.8594
L↓1.4844
-44.4%
Lrecon↓
0.0412
0.0256
L↓0.0156
-37.9%
Lsparse↓
0.2682
0.0209
L↓0.2473
-92.2%
Lalign↓
0.0605
0.0762
L↑0.0157
+26.0%
Lortho↓
0.5330
0.3760
L↓0.1570
-29.5%
Lvalue↓
3.3169
3.4125
L↑0.0956
+2.9%
Table 5: Stage-1 ablation. Deltas are computed as With–No ; losses prefer ↓ and accuracies prefer ↑ .
Figure 7: Swap controllability in 2-hop reasoning. The post-trained AlignSAE achieves 4× higher swap success than the traditional SAE when α≈50 .
Figure 8: Relation-to-slot binding (pre- vs. post-grokking). At steps e2 and e3 , 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 ).
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.
Table 7: Correct swap examples from the evaluation set (Layer 6; α=2 ).
Setting
AlignSAE
Standard SAE
Ratio
1-Hop (6 concepts)
0.173
0.000
∞
2-Hop (20 relations)
0.035
0.004
9.2×
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.
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 α . 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 → birth_date ( α=10 )
27,July,1991 , Bryan Caleb Atkins was born. He is professionally active in Fort Worth, TX…
Steer → major ( α=10 )
French . According to records, Maurice Hudson Tran was born on 28,September,1953…
Steer → work_city ( α=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, α=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 α=10
2.1
1.1
3.9
15
Steered α=100
2.0
1.1
3.9
15
2-Hop
Baseline
4.0
1.2
—
5
Steered α=10
4.0
1.3
2.5
15
Steered α=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.
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 e2 ) and the 2nd hop (Token 2, predicting the final entity e3 ) under step-wise supervision.
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.
Wilson E. Marcílio-Jr, Danilo M. Eler
Adaption Labs, Curitiba, PR, Brazil · São Paulo State University (UNESP), Presidente Prudente, SP, Brazil
Large Language Models (LLMs) have achieved strong complex reasoning capabilities through Chain-of-Thought (CoT) reasoning. However, their reasoning patterns remain too complicated to analyze. While Sparse Autoencoders (SAEs) have emerged as a powerful tool for interpretability, existing approaches predominantly operate at the token level, creating a granularity mismatch when capturing more critical step-level information, such as reasoning direction and semantic transitions. In this work, we propose step-level sparse autoencoder (SSAE), which serves as an analytical tool to disentangle different aspects of LLMs' reasoning steps into sparse features. Specifically, by precisely controlling the sparsity of a step feature conditioned on its context, we form an information bottleneck in step reconstruction, which splits incremental information from background information and disentangles it into several sparsely activated dimensions. Experiments on multiple base models and reasoning tasks show the effectiveness of the extracted features. By linear probing, we can easily predict surface-level information, such as generation length and first token distribution, as well as more complicated properties, such as the correctness and logicality of the step. These observations indicate that LLMs should already at least partly know about these properties during generation, which provides the foundation for the self-verification ability of LLMs. Our code is available at https://github.com/Miaow-Lab/SSAE.
Xuan Yang, Jiayu Liu, Yuhang Lai +3
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
Sparse autoencoders (SAEs) are proposed to extract numerous features from large language model (LLM) representations, yet explaining these features still relies primarily on external observation. This reliance leads to superficial explanations inferred from observed model behavior and computational inefficiency from collecting such behavioral evidence at scale. We introduce SAEVerbalizer, a framework that injects SAE decoder directions into an LLM's representations and fine-tunes the LLM's downstream layers to generate natural-language explanations of the injected features. Once trained, the resulting verbalizer explains SAE features directly from decoder directions, addressing both limitations. Our experiments show that the learned verbalization capability generalizes to unseen features, transfers across separately trained SAE dictionaries, and, with a lightweight adapter, extends to SAE features from different LLMs. Intervention experiments show that injecting multiple directions yields an explanation combining their meanings, while reversing individual directions produces corresponding meaning shifts.
Weihan Meng, Hongzhu Guo, Yi Jing +5
Tsinghua University · Peking University · Fudan University