Search papers, labs, and topics across Lattice.
This paper introduces semi-nonnegative matrix factorization (SNMF) to decompose MLP activations in LLMs into interpretable features, addressing limitations of sparse autoencoders (SAEs) in mechanistic interpretability. SNMF learns sparse linear combinations of co-activated neurons and maps them to activating inputs, enhancing interpretability and enabling causal steering. Experiments on Llama 3.1, Gemma 2, and GPT-2 demonstrate that SNMF-derived features outperform SAEs and a supervised baseline in causal steering while aligning with human-interpretable concepts, revealing a hierarchical structure in MLP activation space.
Forget sparse autoencoders: semi-nonnegative matrix factorization directly dissects MLP activations into human-interpretable features that causally steer LLMs better.
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.
Google Research