Decomposing MLP Activations into Interpretable Features via Semi-Nonnegative Matrix Factorization
This work addresses the challenge of mechanistic interpretability in LLMs by providing a more effective unsupervised method for feature extraction, though it appears incremental as it builds on existing dictionary learning approaches.
The paper tackles the problem of identifying interpretable features in large language models by proposing semi-nonnegative matrix factorization (SNMF) to decompose MLP activations, showing that SNMF-derived features outperform sparse autoencoders and a supervised baseline on causal steering tasks in models like Llama 3.1 and GPT-2.
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.