🤖 machine learning

Tree SAE: Learning Hierarchical Feature Structures in Sparse Autoencoders

The paper introduces Tree SAE, a novel sparse autoencoder model that overcomes the limitations of existing activation-based hierarchical methods by incorporating a reconstruction condition to learn accurate, functional hierarchical feature structures directly from data, thereby outperforming state-of-the-art benchmarks and revealing complex concept hierarchies in large language models.

Original authors: Tue M. Cao, Hoang X. Nhat, Raed Alharbi, My T. Thai

Published 2026-05-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tue M. Cao, Hoang X. Nhat, Raed Alharbi, My T. Thai

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to understand how a giant, complex library (a Large Language Model) organizes its books. You want to find the "shelves" and "categories" inside the library's brain.

For a while, researchers have used a tool called a Sparse Autoencoder (SAE). Think of an SAE as a very efficient librarian who tries to break down the library's contents into simple, single-topic labels. If a book is about "dogs," the librarian tags it "dog." If it's about "cats," they tag it "cat."

However, the real world isn't just a flat list of tags. It's a hierarchy. "Dog" is a type of "Animal." "Golden Retriever" is a type of "Dog." The old librarians (standard SAEs) were terrible at seeing these family trees. They often got confused, mixing up unrelated things or splitting one concept into too many tiny, messy pieces.

The Problem: The "False Friend" Mistake

The old way of finding these family trees relied on a rule called Activation Coverage.

The Rule: "If the 'Dog' tag is on, the 'Animal' tag must also be on."
The Mistake: The paper shows this rule is too loose. Imagine a tag called "Things that exist." It turns on for everything. So, it technically covers "Dog," "Cat," and even "Democrat" (a political party).
The Result: The old system would falsely claim that "Democrat" is a child of "Things that exist" just because they both turn on sometimes. They are semantically unrelated, but the math said they were connected. It's like saying a "Spoon" is a child of "Furniture" just because both are found in a house, even though a spoon isn't a piece of furniture.

The Solution: The "Tree SAE"

The authors built a new librarian called Tree SAE. Instead of just looking at who turns on when, this new librarian checks two things to build a family tree:

The "When" Check (Activation Coverage): Does the parent turn on whenever the child turns on? (Yes, "Animal" is on when "Dog" is on).
The "What" Check (Reconstruction Condition): This is the new, secret sauce. The librarian asks: "If I use the 'Animal' tag and the 'Dog' tag together, do they actually help rebuild the picture of a 'Dog'?"
- If the parent is just a generic "Things that exist" tag, it doesn't help rebuild the specific image of a dog. It fails this test.
- If the parent is a real "Animal" tag, it does help rebuild the dog image. It passes.

By combining these two checks, the Tree SAE builds a strict, accurate family tree where parents and children actually make sense together.

How It Works: The "Privilege" System

The paper describes the Tree SAE as having layers of privilege, like a corporate hierarchy or a tree with roots and branches:

Layer 0 (The Root): The big, general concepts (like "Animal").
Layer 1, 2, etc.: More specific concepts (like "Dog," then "Golden Retriever").

The system is designed so that a specific feature (a child) can only "activate" if its parent is also active. It's like a security system: you can't open the "Golden Retriever" door unless you've already unlocked the "Dog" door, which requires the "Animal" door to be open.

The "Dynamic Re-assignment" Trick

Sometimes, a specific "child" feature gets stuck and stops working (it becomes a "dead feature"). In old systems, this was a permanent problem.
The Tree SAE has a dynamic manager. If a child feature stops working, the manager looks at the "parent" features and asks, "Who needs more help?" It then reassigns the dead child to a parent that actually needs it, ensuring the whole tree stays healthy and active.

What They Found

The paper tested this new Tree SAE against the old methods and found:

Better Family Trees: It found many more correct parent-child pairs and almost none of the "false friend" mistakes.
No More Messy Splitting: It stopped the problem where one concept gets split into too many confusing pieces.
Just as Good at the Basics: It didn't lose its ability to understand the language; it actually reconstructed the original data very well, matching or beating the best existing tools.
Visualizing the Shape: Because the tree is built correctly, researchers can now look at the "shape" of the concepts. They can see how the library organizes "General Adjectives" into specific types like "Subsequent" or "Various," proving the model understands the hierarchy.

In Summary

The paper argues that to truly understand how AI thinks, we need to stop treating its knowledge as a flat list of random tags. We need to build a Tree SAE that forces the AI to organize its knowledge into a proper family tree, checking not just when things happen, but how they fit together to make sense. This results in a much clearer, more accurate map of the AI's internal world.

Technical Summary: Tree SAE: Learning Hierarchical Feature Structures in Sparse Autoencoders

Problem Statement

Sparse Autoencoders (SAEs) are widely used to extract human-interpretable features from opaque language models. However, standard SAEs treat features as an independent set, neglecting the inherent hierarchical structure of real-world concepts. This oversight leads to undesirable phenomena such as feature absorption (where fine-grained features inhibit coarse ones), feature splitting (where a single concept is fragmented across multiple features), and feature composition (where models learn fragmented or overly complex "polysemantic" amalgams).

Recent architectures, such as Matryoshka SAEs and Matching Pursuit SAEs, attempt to address these issues by introducing multi-level reconstruction or conditional orthogonality. They rely on Masked Cosine Similarity (MCS) to identify hierarchical relationships, which is based on the activation coverage condition: a child feature is considered a descendant of a parent if it activates only when the parent activates.

The authors demonstrate that activation coverage alone is insufficient for identifying true hierarchical relationships. They provide evidence that parent and child features can satisfy the activation coverage condition while remaining semantically unrelated (e.g., a "time" concept parent incorrectly paired with a "political" concept child, or a dense "PCA feature" acting as a parent to unrelated concepts). This leads to false positives where the identified hierarchical pairs lack conceptual coherence.

Methodology

1. The Reconstruction Condition

To address the limitations of activation coverage, the authors introduce a novel reconstruction condition. This condition enforces a deeper functional link between hierarchical levels.

Definition: For a parent feature $f_p$ and a child feature $f_c$ , the reconstruction score $S_{res}$ is defined based on the alignment of their decoder vectors ( $d_p, d_c$ ) with the true concept vector of the child ( $d^*_c$ ).
Mechanism: The condition requires that both the parent and child decoder vectors contribute constructively to the reconstruction of the child's concept. Specifically, it ensures the parent captures a broad conceptual direction that the child refines.
Validation: The authors show that while activation coverage may be satisfied by spurious pairs, the reconstruction score effectively filters these out. True hierarchical pairs exhibit high reconstruction scores, whereas unrelated pairs (even with perfect activation coverage) exhibit low scores.

2. Tree SAE Architecture

The authors propose Tree SAE, a model designed to learn hierarchical structures directly within the feature set, rather than inferring them post-hoc.

Tree Structure: The feature set is divided into multiple "privilege layers." A feature at a higher layer can be assigned as a descendant of a feature at any lower layer. An imaginary root node at layer 0 serves as a baseline for features without explicit parents.
Allocation Vector: An allocation vector $a_l$ assigns parent features to child features at each layer $l$ .
Regularized Activation: To enforce the tree structure, the activation of a child feature $f_i$ is regularized to be non-zero only if its assigned parent feature is active. This strictly follows the activation coverage condition during training.
Multi-level Reconstruction Loss: The training objective includes a multi-level reconstruction loss. For a layer $l$ , the loss constrains the cumulative reconstruction of all preceding layers (including $l$ ) to accurately recover the input. This ensures parent vectors capture general concepts and child vectors refine them into specialized concepts.
Dynamic Allocation: To mitigate the issue of "dead features" (features that never activate), the authors introduce a dynamic allocation mechanism. This algorithm reallocates dead child features to parent features based on a "payoff" metric (derived from accumulated training loss), ensuring that child features are assigned to parents that can effectively support their learning.

Key Contributions

Identification of Activation Coverage Limitations: The paper provides concrete evidence that activation coverage alone identifies semantically incoherent hierarchical pairs and proposes a new definition combining activation coverage with the reconstruction condition.
Tree SAE Model: The introduction of a novel SAE architecture that encodes a hierarchical tree structure directly within the feature set, which is dynamically learned during training.
Superior Hierarchical Learning: Demonstration that Tree SAEs learn significantly more coherent parent-child pairs than existing SAEs (including Matryoshka and Matching Pursuit SAEs) while avoiding the pathologies of feature splitting, absorption, and composition.
Geometric Analysis: Showcasing the utility of the explicit tree structure in mapping the geometry of child feature subspaces and understanding how language models decompose concepts into coarse- and fine-grained levels.

Experimental Results

The authors evaluated Tree SAEs on GPT2-small across various sparsity levels ( $L_0$ ) and dictionary sizes, comparing them against Top-k SAEs, Matryoshka SAEs, and Matching Pursuit SAEs.

Hierarchical Metrics: Tree SAEs discovered substantially more hierarchical pairs that satisfied both activation coverage and reconstruction conditions compared to other models. The internal tree structure was found to be slightly more effective at identifying these pairs than post-hoc MCS analysis.
Feature Pathologies: Tree SAEs performed competitively with or outperformed state-of-the-art (SOTA) Matryoshka SAEs in reducing feature splitting, absorption, and composition.
Reconstruction and Interpretability: Tree SAEs maintained competitive reconstruction loss and variance explained, closely reconstructing original latents. In terms of interpretability (measured via AutoInterp), Tree SAEs matched the performance of Matryoshka SAEs.
Scaling: Tree SAEs scaled well with dictionary sizes (6k to 49k), often outperforming Matryoshka SAEs on hierarchical metrics at smaller scales and maintaining comparable performance at larger scales.
Child Feature Diversity: The learned child features exhibited low co-activation rates (<3.0%), indicating that they represent distinct aspects of the parent concept rather than redundant information.

Significance and Claims

The paper claims that Tree SAE represents a significant step forward in understanding the internal representations of large language models. By moving beyond independent feature sets to explicitly model hierarchical structures, the authors argue that Tree SAEs:

Mitigate known hierarchical feature pathologies (splitting, absorption, composition).
Provide a more accurate framework for identifying true parent-child concept relationships, avoiding the false positives common in previous methods.
Enable the geometric analysis of hierarchical concept subspaces, offering new insights into how models decompose general concepts into specialized sub-concepts.

The authors note limitations, specifically that the rigid tree structure can lead to a higher number of dead features compared to other SAEs, though dynamic allocation helps mitigate this. They suggest that future work might explore soft-gating mechanisms to further improve feature utilization.