Disentangling Dense Embeddings with Sparse Autoencoders

TL;DR

This work introduces sparse autoencoders (SAEs) as a means to disentangle semantic content in dense text embeddings derived from large language models, addressing interpretability and control gaps in semantic search. By training SAEs on embeddings from hundreds of thousands of abstracts across astronomy and computer science, the authors uncover interpretable features, quantify interpretability with automated LLM-based labelling, and reveal a hierarchical organization via feature families built with a graph-based clustering approach. They demonstrate practical utility by steering semantic search through feature- or family-level interventions, showing improved specificity over traditional query rewriting, and provide an open-source toolkit including SAEs, embeddings, and a web app for exploration. The findings establish scaling laws for SAE performance, illuminate the structure of learned features, and offer a replicable pathway to more interpretable and controllable semantic spaces in NLP tasks.

Abstract

Sparse autoencoders (SAEs) have shown promise in extracting interpretable features from complex neural networks. We present one of the first applications of SAEs to dense text embeddings from large language models, demonstrating their effectiveness in disentangling semantic concepts. By training SAEs on embeddings of over 420,000 scientific paper abstracts from computer science and astronomy, we show that the resulting sparse representations maintain semantic fidelity while offering interpretability. We analyse these learned features, exploring their behaviour across different model capacities and introducing a novel method for identifying ``feature families'' that represent related concepts at varying levels of abstraction. To demonstrate the practical utility of our approach, we show how these interpretable features can be used to precisely steer semantic search, allowing for fine-grained control over query semantics. This work bridges the gap between the semantic richness of dense embeddings and the interpretability of sparse representations. We open source our embeddings, trained sparse autoencoders, and interpreted features, as well as a web app for exploring them.

Figures (27)

• Figure 1: Training and feature labelling process for our sparse autoencoder (SAE). The SAE is trained to minimise reconstruction loss on embeddings from astronomy and computer science paper abstracts. Each feature corresponds to a column in the decoder matrix, representing a direction in embedding space. Feature interpretation involves two steps: (1) An Interpreter language model identifies topics present in text that activates each feature but absent in non-activating text. (2) A separate Predictor language model assesses feature interpretability by stating its confidence that the feature will activate on unseen text, with confidence correlated with ground truth activations to quantify interpretability.
• Figure 2: Scaling laws for sparse autoencoder performance. Left: Normalised mean squared error (MSE) as a function of the number of total latents $n$ for different values of active latents $k$. The power-law scaling is evident for each $k$. Right: Reconstruction loss as a function of compute (FLOPs) for different $k$ values, demonstrating the compute-optimal model size scaling.
• Figure 3: Pearson correlations between the ground-truth and predicted feature activation, using GPT-4o as the Interpreter and GPT-4o-mini as the Predictor.
• Figure 4: Sample feature family from cs.LG; arrows represent $C_{ij}^{norm} > 0.1$, with size $\propto$$C_{ij}^{norm}$.
• Figure 5: Co-occurrence matrix $C$ organised by a subset of 5 feature families each. Features in families are ordered by firing density, and the right-most feature is the parent. The un-filled block structure reflects the hierarchical nature of the feature family: all children co-occur with the parent, but few children fire with each other. Visually, this supports our clustering approach.