Constructing Interpretable Features from Compositional Neuron Groups

TL;DR

This paper introduces a semi-nonnegative matrix factorization method to extract interpretable, sparse features from language model activations, improving causal interpretability and revealing hierarchical concept structures.

Contribution

It proposes a novel SNMF-based approach for unsupervised feature extraction that outperforms existing methods like SAEs in interpretability and causal evaluation.

Findings

SNMF features outperform SAEs in causal steering tasks.

Features align with human-interpretable concepts.

Neuron combinations are reused across related features, indicating hierarchy.

Abstract

A central goal for mechanistic interpretability has been to identify the right units of analysis in large language models (LLMs) that causally explain their outputs. While early work focused on individual neurons, evidence that neurons often encode multiple concepts has motivated a shift toward analyzing directions in activation space. A key question is how to find directions that capture interpretable features in an unsupervised manner. Current methods rely on dictionary learning with sparse autoencoders (SAEs), commonly trained over residual stream activations to learn directions from scratch. However, SAEs often struggle in causal evaluations and lack intrinsic interpretability, as their learning is not explicitly tied to the computations of the model. Here, we tackle these limitations by directly decomposing MLP activations with semi-nonnegative matrix factorization (SNMF), such that the learned features are (a) sparse linear combinations of co-activated neurons, and (b) mapped to their activating inputs, making them directly interpretable. Experiments on Llama 3.1, Gemma 2 and GPT-2 show that SNMF derived features outperform SAEs and a strong supervised baseline (difference-in-means) on causal steering, while aligning with human-interpretable concepts. Further analysis reveals that specific neuron combinations are reused across semantically-related features, exposing a hierarchical structure in the MLP's activation space. Together, these results position SNMF as a simple and effective tool for identifying interpretable features and dissecting concept representations in LLMs.