Tree SAE: Learning Hierarchical Feature Structures in Sparse Autoencoders

TL;DR

Tree SAE is a novel autoencoder model that learns hierarchical feature structures by combining activation and reconstruction constraints, improving the understanding of complex data hierarchies.

Contribution

The paper introduces Tree SAE, a new method that effectively captures hierarchical relationships in features by integrating activation and reconstruction conditions.

Findings

Tree SAE outperforms existing SAEs in learning hierarchical feature pairs.

Tree SAE maintains competitive performance on key benchmarks.

Demonstrates utility in analyzing hierarchical concepts in language models.

Abstract

Learning hierarchical features in Sparse Autoencoders (SAEs) is essential for capturing the structured nature of real-world data and mitigating issues like feature absorption or splitting. Existing works attempt to identify hierarchical relationships within independent feature sets by relying on activation coverage, the assumption that child feature should only activate when its parent feature activates. However, we demonstrate that this condition alone is insufficient; that is, it often produces false positives where parent and child concepts are semantically unrelated. To address this, we introduce a novel reconstruction condition that enforces a deeper functional link between hierarchical levels. By combining both activation and reconstruction constraints, we propose the Tree SAE, a model designed to learn hierarchical structures directly from within the feature set. Our results demonstrate that Tree SAEs significantly surpass the existing SAEs at learning hierarchical pairs while maintaining competitive performance to the state-of-the-art on several key benchmarks. Finally, we demonstrate the practical utility of our Tree SAE in mapping the geometry of child feature subspaces and uncovering the complex hierarchical concept structures encoded within large language models.