TaxoSAE: Rethinking Sparse Autoencoders through Taxonomic Logic

1 Introduction

• Earlier concept-alignment methods bind concepts to SAE features, while contrastive activation differences identify semantic features in unsupervised SAEs.These approaches support interpretation and targeted feature interventions, but do not themselves encode taxonomic dependencies.Source: §1 · p. 1
• Flat SAE features lack taxonomic consistency for readout and cannot specify an edit’s scale or which dependent information should remain unchanged.For example, replacing one cell type with another may preserve its hypernym, while replacing it with a different kind of cell changes the hypernym too.Source: §1 · pp. 1–2
• TaxoSAE mirrors taxonomy levels with stacked SAEs and trains each concept to activate whenever one of its hyponyms activates.The approach requires explicit data points only for terminal concepts, given the taxonomy.Source: §1 · p. 2
• +31.3 percentage points in readout accuracy on questions about abstract concepts never seen in training, alongside gains in control and unseen-concept hypernym readout.TaxoSAE uses 1.0–1.7 slots per abstract concept; swap success improves by +25.5 percentage points, and unseen-concept hypernym readout improves by up to +37.9 percentage points.Source: §1 · p. 2

2 Related Work

• Supervised concept-alignment SAEs bind concepts to dedicated slots, whereas TaxoSAE organizes the slots into levels linked by is-a relations.Unsupervised SAEs require post-hoc feature selection, and a single concept may be split across features or absorbed into more specific ones.Source: §2 · p. 2
• Prior analyses find taxonomic structure in LLM representation geometry but focus on discovering structure rather than reading out and controlling concepts consistently across levels.TaxoSAE makes the taxonomy explicit in the feature space so concepts and hypernyms can be read out and errors edited.Source: §2 · pp. 2–3
• Figure 1 depicts taxonomy edges as is-a links from hypernyms to hyponyms and illustrates an Animal taxonomy example.The figure also indicates that TaxoSAE uses one SAE per taxonomy level.Source: Figure 1 · p. 3
• Taxonomic alignment encodes the rule “if a hyponym holds, its hypernym holds” as a soft constraint across stacked SAEs.This follows the training-objective approach to encoding logic rules in neural networks.Source: §2 · p. 3

3 Task Formulation

• A taxonomy is a directed acyclic graph in which each node has at most one hypernym; roots have no parents, terminal nodes have no children, and other nodes are intermediate.Levels are indexed from the bottom: L1 contains terminal nodes and LK contains roots.Source: §3 · p. 3
• Factual QA is formulated over triples (e, r, v), where a question specifies an entity and relation and the model returns a value.The taxonomic element t is the relation for relation taxonomies and the entity for entity taxonomies; the other element is complementary c.Source: §3 · p. 3

4.1 T axoSAE Architecture

• One SAE per taxonomy level maps the lower level’s reconstruction to its hypernyms, abstracting a question from its terminal concept upward.The first SAE reads the frozen LLM’s last-token hidden state; each higher SAE reads the reconstruction below.Source: §4.1 · p. 3
• Taxonomic slots, complement slots, and free slots represent taxonomy nodes, complementary elements, and remaining reconstruction information, respectively.The value head predicts answers from taxonomic and complementary slots, tying each slot to facts of its node.Source: §4.1 · p. 4

4.2 T axoSAE Training

• The jointly trained objective combines standard SAE reconstruction and sparsity losses with taxonomic and complement alignment, decorrelation penalties, and value prediction.Pre-training first learns a plain dictionary with the SAE loss alone; post-training optimizes the full objective jointly.Source: §4.2 · pp. 4–5
• Non-terminal hypernym targets are derived by summing the lower-level probabilities of their hyponyms, then selecting the hypernym with the largest total.When the lower SAE reads its node correctly, this target is that node’s hypernym.Source: §4.2 · pp. 4–5

5.1 Datasets and Baselines

• Bio contains 18 terminal, 6 intermediate, and 3 top-level relation nodes; Animal contains 18 terminal, 6 intermediate, and 2 top-level animal nodes.Training and out-of-distribution questions differ by templates for Bio and synonyms for Animal.Source: §5.1 · p. 5
• Baselines include random feature assignment, difference-in-means feature selection, and a hybrid using TaxoSAE terminal slots with selected unsupervised features at higher levels.The unsupervised SAEs use reconstruction and L1 losses.Source: §5.1 · p. 5

5.2 Implementation

• The base model is fine-tuned GPT-2 small, and TaxoSAE uses three stacked SAEs trained on terminal-level questions only.The hidden state is taken at the final prompt token from a middle layer; Bio is also evaluated with Qwen2.5-0.5B.Source: §5.2 · p. 5

5.3 Metrics

• Readout checks whether the most active unit matches a queried node and its hypernyms, while control tests slot swaps and selective deactivation.Activation concentration is measured with effective feature count and the strongest unit’s share; deactivation is evaluated on the node’s own versus other questions.Source: §5.3 · pp. 5–6
• Table 1 reports that TaxoSAE achieves the best readout and control on both Bio and Animal datasets.The table defines L1/L2/L3 as terminal/intermediate/top levels and reports level-r readout on level-q questions.Source: Table 1 · p. 6

6.1 Taxonomy Readout

• 31.3 points: TaxoSAE gains over the strongest baseline in L2→L3 hypernym readout on non-terminal questions, and gains up to 25.4 points on L2→L2 and L3→L3 readout.It is trained only on terminal-level questions; complementary slots read out almost perfectly.Source: §6.1 · p. 7
• 1.0–1.7 EffFeat and 89.0–100.0% Top1C: TaxoSAE concentrates abstract concepts in about one slot, while unsupervised SAE activations fragment across features.Figure 4 shows Bio activations and binding matrices; Figure 5 compares fragmentation and concentration across levels on Animal and Bio.Source: §6.1 · p. 7, Figure 4 · p. 7, Figure 5 · p. 7

6.2 Taxonomy Controllability

• 25.5 points: TaxoSAE improves intermediate-level target-only swap coverage, while target hypernyms are read out in 93.2–100.0% of swaps, 28.9–37.3 points above baselines.The comparison is for Bio swaps at each system’s best strength.Source: Table 2 · p. 7, §6.2 · p. 8
• 0.00–0.02 target accuracy versus 0.91–1.00 accuracy on other elements: TaxoSAE suppresses target answers while preserving non-target answers across both datasets.DiffMean suppression also disrupts other answers, especially at the intermediate level.Source: §6.2 · p. 8, Figure 6 · p. 8

6.3 Ablation Study

• Removing taxonomic alignment collapses readout and swap success to random-baseline levels; removing other losses mainly reduces swap success.Orthogonality-loss ablations have smaller swap-success drops, while readout remains comparable; some concentration metrics improve.Source: §6.3 · p. 8, Figure 7 · p. 8

• Flat SAE feature spaces leave readout without taxonomic consistency and control without granularity.Source: §7 · p. 8
• TaxoSAE combines stacked SAEs with a taxonomic alignment loss to encode is-a logic in the feature space.Source: §7 · p. 8
• On two factual QA taxonomies, TaxoSAE improves taxonomy readout and control over baselines and reads unseen concepts at their hypernyms without binding them.The authors identify larger models and logical relations beyond taxonomies as future work.Source: §7 · p. 8

Limitations

• Experiments are in English; transfer to other languages, where taxonomic knowledge may be expressed and encoded differently, is untested.Source: Limitations · p. 9
• The taxonomies have three levels and at most 18 terminal nodes, limiting the evaluated taxonomy complexity.Source: Limitations · p. 9

A Datasets

• Both datasets use controlled facts with known answers; TaxoSAE trains on terminal-level questions, while intermediate- and top-level questions are readout tests.Bio covers biographical relations and Animal covers animals; both use complementary elements across a three-level taxonomy.Source: A Datasets · p. 11

B Results on Qwen2.5-0.5B

• On Bio with Qwen2.5-0.5B, TaxoSAE reads non-terminal concepts best, scoring 15.6–33.8 points above the strongest baseline, and represents each with 1.0–1.3 effective features versus more than 100 for baselines.Source: B Results on Qwen2.5-0.5B · p. 11, Table 4 · p. 12
• 0.3% own-question accuracy and at least 99.8% other-question accuracy: deactivating a TaxoSAE slot selectively removes the target concept’s answers.DiffMean cannot reduce a node’s own accuracy below 5% at the tested strengths.Source: B Results on Qwen2.5-0.5B · p. 11, Table 4 · p. 12
• 60.5% and 13.7%: TaxoSAE’s absolute swap success is lower than on GPT-2, which the authors say may partly reflect the untuned layer choice.Source: B Results on Qwen2.5-0.5B · p. 11, Table 4 · p. 12
• 76.2% L1 swap and 36.5% L2 swap at layers 16–18: control improves over layer 12 by 15.7 and 22.8 points, then declines at later layers.Readout remains high at every tested layer.Source: B Results on Qwen2.5-0.5B · pp. 11–12

C Unseen Terminal Concepts

• On unseen terminals, Animal L2/L3 accuracy was 63.9%/76.1% (vs. 26.0%/56.8% for DiffMean); Bio was 91.4%/94.8% (vs. 75.9%/69.8% for the strongest baselines).Animal errors were concentrated on a few concepts; 12 of 24 unseen animals were classified correctly on at least 99% of their questions.Source: C Unseen Terminal Concepts · pp. 12–13, Table 7 · p. 15

D Unseen Phrasings

• T axoSAE retains its advantage over baselines on unseen rephrasings or synonyms, with near-perfect Bio accuracy and a clear Animal lead at non-terminal levels.The figure reports diagonal accuracy for seen templates and unseen rephrasings or synonyms across L1–L3.Source: D Unseen Phrasings · p. 14, Figure 10 · p. 15

E Complement Slot Readout

• Complement slots read the question’s complementary person or property almost perfectly on terminal questions, including unseen templates, names, and synonyms.The reported evaluation uses macro accuracy over complementary elements.Source: E Complement Slot Readout · p. 15, Table 8 · p. 15

F.1 Loss Ablation

• Removing taxonomic–free orthogonality improves some concentration metrics but lowers intermediate-level swap coverage from 41.8% to 24.2%.Removing taxonomic alignment collapses readout, while removing other losses changes readout much less.Source: §F.1 · p. 15

F.2 Value head Ablation

• Replacing the bilinear head with a per-slot head lowers intermediate swap success from 41.8% to 38.3% on Bio and 27.6% to 22.1% on Animal.Readout, concentration, and selectivity barely change; the intermediate-level drops exceed variation across three training seeds.Source: §F.2 · p. 15, Table 11 · p. 16

G Flat AlignSAE Baseline

• On Bio, flat AlignSAE achieves 5.5% terminal swap success versus 84.7% for T axoSAE, and 18.3% versus 41.8% on intermediate swaps.Both systems read terminal-level questions perfectly at every level, but the flat SAE leaves higher-level slots unchanged after edits.Source: G Flat AlignSAE Baseline · pp. 16–17, Table 12 · p. 17

H Swap Success Results

• After swaps, T axoSAE reads target hypernyms in 93.2–100% of cases, compared with at most 71.1% for baselines.After terminal swaps, intermediate readout is 100% on both datasets; after intermediate swaps, top-level readout is 100% on both.Source: H Swap Success Results · p. 17, Table 13 · p. 18

J Baseline Implementation Details

• Baseline dictionaries are trained unsupervised, while post-hoc feature selection uses labels; non-terminal labels are derived from terminal labels through the taxonomy.The max-margin difference-in-means selection favors features that respond more to a concept than to other elements at the same level.Source: J Baseline Implementation Details · p. 18

I Selectivity Matrices with Concept Labels

• Figure 12’s rows deactivate an element’s unit, and its columns report exact-match accuracy for each element.The figure labels the concepts represented by its rows and columns.Source: Figure 12 · p. 19

K.1 Base Model Training

• Training combines terminal factual QA, hypernymy QA, and plain-text descriptions, but includes no non-terminal factual QA.Hypernymy QA supplies taxonomic structure without supplying non-terminal factual answers.Source: §K.1 · p. 19

K.2 T axoSAE Training

• TaxoSAE trains in two stages: Stage 1 uses reconstruction and sparsity losses, then Stage 2 enables all losses and derives non-terminal alignment targets from terminal predictions through the taxonomy.Only terminal-level taxonomic alignment uses ground-truth labels.Source: §K.2 · p. 20

L Layer-wise Performance

• On Bio, layer 6 has the best terminal swap at 84.7%; on Animal, layer 9 has the best terminal readout at 84.3% and terminal swap at 43.6%.Animal layer 10 reaches the highest intermediate-level swap, 54.5%, versus 27.6% at layer 9.Source: L Layer-wise Performance · p. 21, Table 19 · p. 21, Table 20 · p. 22

M The Usage of AI Tools

• The authors state that LLMs were used only for auxiliary manuscript-language refinement, not research ideation, hypothesis formulation, or experimental design.They attribute conceptual contributions, experimental designs, and method design to the authors.Source: M The Usage of AI Tools · p. 21