Interpretable Embeddings with Sparse Autoencoders: A Data Analysis Toolkit

AI summary

Overview

Research area: Mechanistic interpretability and data-centric interpretability — using sparse autoencoders (SAEs) trained on LLM hidden states to build interpretable text embeddings for large-scale corpus analysis.

Technical level: Intermediate. The paper assumes familiarity with LLM hidden states and embedding models, but the core pipeline (feed text to a reader LLM, capture SAE activations, max-pool across tokens) is described concretely enough for readers with basic NLP background.

Scope: The paper proposes "SAE embeddings" — representations whose individual dimensions map to human-interpretable concepts — and evaluates them across four data analysis tasks (dataset diffing, correlation discovery, clustering, retrieval) plus two model-behavior case studies.

What This Paper Is About

Analyzing large text corpora is important for finding undesirable model behaviors and dataset biases, but existing tools force a trade-off: LLM-based labeling is expensive and prompt-sensitive, while dense embeddings are fast but opaque and offer no control over which properties you examine. The authors show that SAEs trained on LLM hidden states can produce embeddings whose dimensions correspond to interpretable concepts, effectively labeling text with thousands of concepts at once. The goal is a toolkit that is cheaper and more reliable than LLM baselines while retaining the controllability that dense embeddings lack.

Key Contributions

1. A method for interpretable text embeddings: Each document is passed through a "reader LLM," the SAE activations are max-pooled across tokens, and each resulting dimension maps to a human-understandable concept label — yielding a single sparse embedding per document.

2. Demonstration across four data analysis tasks: dataset diffing (Section 4.1), finding correlations between arbitrary concepts (4.2), clustering (4.3), and property-based retrieval (4.4), each validated on ground-truth datasets before being applied to unlabeled data.

3. Cost and reliability comparisons against LLM and dense-embedding baselines: The authors report that SAE embeddings uncover bigger differences at 2–8x lower cost than frontier LLMs and identify biases more reliably.

4. Two applied case studies on real model behavior: tracking how OpenAI model behavior has changed across generations, and discovering a spurious prompt–response correlation in Tulu-3's post-training data (the phrase "I hope it is correct").

Main Findings

• Dataset diffing beats LLM baselines in magnitude and cost: When diffing model outputs, SAE hypotheses showed a higher average judge-verified frequency difference per hypothesis than the LLM summarize (LLM-S) and cluster (LLM-C) baselines (Figure 2). Generating hypotheses with pure LLMs was reported as 2–8x more expensive than SAEs; SAEs are especially cost-effective in multi-model comparisons because the embeddings can be reused, whereas baselines must reprocess responses for each comparison.

• Ground-truth diffing validation: On a movie description dataset with labeled genres and a model responses dataset with tone shifts, the average surface similarity of the top five latents to ground truth (scored 1 = same, 0.5 = related, 0 = unrelated) was 0.75 for movies and 0.8 for tones.

• Specific model differences surfaced: Comparing model responses, the authors find Grok-4 clarifies ambiguities more often than nine other frontier models; on the Grok-4 diff, the top hypotheses were "clarifies ambiguities" (+20%) and "invites interaction" (+46%). LLaVA-Next vs. Vicuna-7B showed nonsensical artifacts like "everybody" or "nobody" (+8.6%) and sarcastic language (+5%). Evaluation-prompt vs. no-prompt showed structured thinking (+19%) and formal definitions in academic tone (+18%).

• Correlations recover synthetic signals that LLMs miss: Ten injected synthetic correlations were placed in a background corpus of 990 Pile texts (Croatian text with emojis; baseball rules with slang; conservative economic opinions in academic tone). The SAE method recovered these small correlations, while the LLM baseline failed to recover them reliably and could fail when the dataset was shuffled, even at temperature = 0 (Figure 3).

• Correlations have better signal-to-noise on real data: On 5k internet comments from CivilComments and a 5k Pile sample, discovered latent pairs (e.g. NPMI_SAE > 0.6, sim < 0.2) generally had higher verified NPMI than pairs raised by the LLM and correlated topic model baselines (Figure 4). Example real-world findings: offensive-language latents co-occur with race, gender, and religion latents on CivilComments; on the Pile, Q&A latents co-occur with software latents (StackExchange-style discussions) and biographical latents co-occur with category-related latents (Wikipedia articles with category metadata).

• Clustering is controllable along chosen axes: On 960 synthetic news paragraphs with four axes of variation (topic, sentiment, temporal framing, writing style), SAE clusters tracked each individual axis more precisely than baselines, which gave topic clusters. On 1k GSM8k solutions, filtering to reasoning-related latents produced clusters of different reasoning approaches, while dense and instruction-tuned embeddings clustered by math problem content instead (Figure 5). SAE clusters were found to have comparable accuracy with dense embedding clusters on other datasets, though in dense embedding space SAE clusters can look "random" (measured via conductance z-score).

• Retrieval works best for model-related data: Across 6 datasets (10k texts each) — ChatbotArena prompts & responses, DeepSeek-R1 reasoning traces, Pile documents, arXiv q-bio abstracts, and Reddit short stories — SAE embeddings generally outperformed or matched baselines on MAP and mean precision@50. Performance was strongest on model-related data (chat responses, reasoning traces, the Pile), which is notably similar to the SAE's training dataset (LMSYS-Chat-1M). Removing the LLM latent reranking step reduced cost but slightly degraded performance. Aggregating the strongest baseline (OpenAI+LLM) with the SAE beat any individual method.

• SAEs capture implicit rather than semantic properties: For the query "model stuck in repetitive loop," dense embeddings returned a document about repetitive loops ("The context memory is getting corrupted"), whereas SAE embeddings returned a document with repetitive loops ("de la peur et de la peur et").

• OpenAI model generations show emerging trends: Comparing five OpenAI models from GPT-3.5-turbo to GPT-5 on 1k sampled general chat prompts, later generations increasingly gave nuanced explanations acknowledging trade-offs or critiques. Starting from GPT-4.1, models began offering personalized follow-ups. A correlated latent pair ("Roleplay scenarios," "personification of objects") was verified by generating 185 role-play prompts with GPT-4o; models increasingly personify objects during role-plays, with GPT-5 almost always doing so.

• Tulu-3 learned a spurious trigger phrase: On a 10k subset of tulu-3-sft-mixture, "math"/"list"/"LaTeX" features in prompts correlated with "hope" features in responses — math prompts whose assistant response contains "I hope it is correct." Tulu-3 does produce this phrase (Llama-3.1-8B-Instruct never does), and a multi-part problem containing a character triggered it even in coding prompts. The authors note the LLM baseline did not find this correlation.

• Methods and baselines used: Goodfire's SAEs trained on layer 50 hidden states of Llama 3.3 70B using LMSYS-Chat-1M, with dictionary size d_SAE = 65536; 61521 existing latent descriptions were reused, with occasional relabeling. Gemini 2.5 Flash was the primary LLM for latent labeling and judging, and OpenAI's text-embedding-3-large was used for dense embeddings and similarity search. SAEs were also benchmarked with Llama 70B and 8B variants. (The reader model size ablation is in Appendix J.)

Methodology in Plain English

The pipeline works as follows. A sparse autoencoder is a network trained to decompose an LLM's internal activation on a single token into a large, sparse set of directions; each direction ("latent") tends to correspond to one interpretable concept. The authors take a pretrained SAE and give each latent a text label by showing an LLM ten sample passages where it fires and ten where it doesn't, then asking for a concise description.

To embed a document, they feed it to a "reader LLM," record the SAE activations at each token, and take the maximum activation across tokens for each latent. The result is a single sparse vector where every dimension has a human-readable label.

This embedding can then be used two ways. As a data labeler, each latent is binarized into a yes/no label for whether the document contains that concept; the authors compare label frequencies across datasets (diffing) or compute co-occurrence between latent pairs (correlations, measured with normalized pointwise mutual information, NPMI, filtered to high-NPMI/low-semantic-similarity pairs to avoid obvious ones). As a controllable embedding, they filter the vector down to only the latents whose labels match a natural-language query, then cluster documents (spectral clustering on a Jaccard similarity matrix) or rank documents for retrieval (weighted sum of filtered latent activations).

Each task was first validated on data with known ground truth — labeled movie genres, synthetic tone shifts, injected synthetic correlations, synthetic news paragraphs — before being applied to unlabeled corpora and compared against LLM-prompting baselines, dense embeddings (OpenAI text-embedding-3-large), instruction-tuned embeddings (Instructor-Large), topic models, and BM25 with LLM query expansion.

Why This Matters

Impact on research: The paper argues that interpretability has focused overwhelmingly on model internals, and that interpreting models through their data is an understudied direction. SAE embeddings provide a reusable, interpretable representation that can support auditing of training data and model outputs without predefined labels or priors — the Tulu-3 case study shows a dataset-auditing insight turning into a testable hypothesis about model behavior.

Real-world applications:

• Dataset auditing and bias detection: surfacing spurious prompt–response correlations in fine-tuning data (as with Tulu-3) and demographic correlations in comment corpora like CivilComments.
• Model comparison and release evaluation: cheaply characterizing how a new model's outputs differ from nine frontier models, including granular behavioral traits like clarifying ambiguities or inviting interaction.
• Exploratory corpus analysis: clustering documents along a chosen axis (e.g. reasoning style in GSM8k solutions) rather than only by topic, and retrieving texts by implicit properties such as tone, hedging, or sycophancy.
• Training-data provenance: surfacing source artifacts, as when Q&A latents co-occurring with software latents traced back to StackExchange-style content and biographical/category latents to Wikipedia metadata.

Industry relevance: The cost argument is central — SAE hypotheses were reported at 2–8x lower cost than LLM-based hypothesis generation, and the embeddings are computed once and reused across comparisons, which matters for teams repeatedly diffing many model versions or large corpora. The controllability of filtering by feature also gives practitioners a knob that dense embedding pipelines do not offer.

Future Directions

• Reducing SAE-specific weaknesses: The authors flag feature absorption as a limitation inherited from prior SAE research, and note that results are sensitive to which latents an SAE learns, which depends on its training data and shapes the available hypothesis space.
• Improving the embedding itself: They state that many design choices — such as how latents are aggregated and which metrics are used — can be refined.
• Better and more varied SAEs: Different SAE sizes and domain-specific SAEs are named as exciting directions, which would address the finding that retrieval quality was best on data distributions similar to the SAE's feature-labeling data.
• Reducing manual effort in correlation discovery: On the Pile, the authors note that some manual effort is required because of the large number of possible latent pairs, suggesting an open problem in automating the triage of candidate correlations.

Target Audience

Researchers and engineers working on mechanistic interpretability, model auditing, dataset curation, and training-data analysis who want interpretable alternatives to LLM-as-annotator pipelines and dense embedding similarity. It is most useful for practitioners building tooling around large text corpora — for example, teams comparing model outputs, screening fine-tuning data, or doing exploratory cluster and retrieval analysis — and for interpretability researchers interested in the data-centric framing. Readers should be comfortable with embedding-based retrieval and the basic mechanics of LLM internals; the paper notes that many methodological details are deferred to its appendices.

Authors’ abstract

Analyzing large-scale text corpora is a core challenge in machine learning, crucial for tasks like identifying undesirable model behaviors or biases in training data. Current methods often rely on costly LLM-based techniques (e.g. annotating dataset differences) or dense embedding models (e.g. for clustering), which lack control over the properties of interest. We propose using sparse autoencoders (SAEs) to create SAE embeddings: representations whose dimensions map to interpretable concepts. Through four data analysis tasks, we show that SAE embeddings are more cost-effective and reliable than LLMs and more controllable than dense embeddings. Using the large hypothesis space of SAEs, we can uncover insights such as (1) semantic differences between datasets and (2) unexpected concept correlations in documents. For instance, by comparing model responses, we find that Grok-4 clarifies ambiguities more often than nine other frontier models. Relative to LLMs, SAE embeddings uncover bigger differences at 2-8x lower cost and identify biases more reliably. Additionally, SAE embeddings are controllable: by filtering concepts, we can (3) cluster documents along axes of interest and (4) outperform dense embeddings on property-based retrieval. Using SAE embeddings, we study model behavior with two case studies: investigating how OpenAI model behavior has changed over time and finding "trigger" phrases learned by Tulu-3 (Lambert et al., 2024) from its training data. These results position SAEs as a versatile tool for unstructured data analysis and highlight the neglected importance of interpreting models through their data.