Adaptive graph sparse autoencoders improve language model feature learning
When Trees Are Not Enough: Learning Mixed-Topology Feature Graphs with Adaptive Graph Sparse Autoencoders
Machine LearningComputation and Language
Summary
Understanding how large language models represent information is tricky because their features often depend on multiple other features, not just one. Previous methods assumed simple tree-like feature relationships that limited discovery. This paper introduces a new method that learns more complex graphs where features can have multiple parents, allowing better organization and meaning extraction. The approach uses this learned structure as a training signal to improve feature discovery and reliability in interpreting language model activations.
What this means in practice
- •For machine learning engineers: Improve interpretability tools by organizing learned features in language models with flexible multi-parent graphs for better semantic insight.
- •For ai model developers: Enhance unsupervised training of feature extractors to yield more reliable and causally controllable components from large language model activations.
Authors
Xiaozuo Shen, Yifei Cai, Tian Tan, Rui Ning, Chunsheng Xin, Hongyi Wu
Abstract
Sparse autoencoders (SAEs) expose interpretable features in large language model activations, yet existing structured SAEs impose single-parent trees or forests, while post-hoc graphs permit multiple parents but neither guide feature learning nor ensure reliable relation recovery. We introduce the Adaptive Graph Sparse Autoencoder (AG-SAE), a structure-guided training paradigm that treats each feature's complete parent set as an atomic structural hypothesis and lets evidence select zero, one, or multiple parents. By competing complete parent sets against null, subset, and alternative explanations, AG-SAE identifies jointly necessary multi-parent relations while rejecting redundant or spurious alternatives and verifying that each child contributes beyond its parents. The induced topology over SAE features then defines a differentiable structural loss that guides SAE training, while topology-guided refinement mitigates feature absorption and uses persistent reconstruction gaps exposed by the learned structure to initialize new features. The entire graph is then induced again from the revised dictionary by reassessing every feature's complete parent set, closing the dictionary-graph self-consistency cycle. Experiments demonstrate exact mixed-topology recovery in a controlled toy model, greater relational reliability and semantic validity than structured and post-hoc baselines on real LLM activations, and stronger feature-level causal interventions than conventional SAE features. AG-SAE thereby turns recovered mixed-topology feature structure into an unsupervised training signal that improves the dictionary, enables reliable feature organization beyond the topological limitations of trees, and exhibits stronger causal control beyond reconstruction.