TaxoSAE improves language model feature organization with taxonomy logic
TaxoSAE: Rethinking Sparse Autoencoders through Taxonomic Logic
Computation and Language
Summary
Modern language models try to break down their knowledge into understandable pieces called features, but these features often don’t follow how real-world ideas are organized. The authors propose TaxoSAE, which arranges these features to follow a hierarchy like categories and subcategories, similar to a family tree. This approach helps read and modify the knowledge more accurately and even understand new concepts by looking at their broader categories.
What this means in practice
- •For natural language processing teams: Enhance language model interpretability and control by encoding hierarchical knowledge, enabling more precise updates and accurate querying of conceptual information.
- •For knowledge graph engineers: Integrate taxonomic alignment in autoencoder features to improve readout and manipulation of hierarchical data in knowledge bases.
Authors
Xinyu Guo, Minglai Yang, Clayton T. Morrison, Eduardo Blanco, Mihai Surdeanu
Abstract
Recent Sparse Autoencoders (SAEs) decompose the representations of large language models (LLMs) into interpretable features, enabling precise readout and control of encoded knowledge. However, existing SAEs treat each feature as a standalone unit, whereas real-world knowledge is often taxonomically organized. As a result, readout carries no guarantee of taxonomic consistency, and control lacks the granularity to specify the scope of an edit. We propose TaxoSAE, which explicitly encodes taxonomical information, built on is-a (hypernym-hyponym) relations, into the SAE feature space. TaxoSAE mirrors a taxonomy with a stack of SAEs, one per taxonomy level, and introduces a taxonomic alignment loss that trains each concept to be active whenever one of its hyponyms is. On two factual QA taxonomies, one over biographical relations and one over animals, TaxoSAE enables (1) more accurate readout through the taxonomy with concentrated activations, e.g., +31.3 points in reading out the hypernyms of abstract concepts never queried in training; (2) more precise control, with +25.5 points in swap success, while suppressing a concept removes only its own answers (0-2% accuracy vs. 91-100% on other concepts); and, more importantly, (3) readout of unseen concepts through their hypernyms, with up to +37.9 points in hypernym readout accuracy.