💬 NLP

Vector Quantized Latent Concepts: A Scalable Alternative to Clustering-Based Concept Discovery

The paper proposes Vector Quantized Latent Concepts (VQLC), a scalable framework that learns interpretable, discrete latent concepts from frozen LLM hidden states by offering a computational efficiency comparable to K-Means while achieving the semantic coherence of hierarchical clustering.

Original authors: Xuemin Yu, Ankur Garg, Samira Ebrahimi Kahou, Hassan Sajjad

Published 2026-06-11

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Xuemin Yu, Ankur Garg, Samira Ebrahimi Kahou, Hassan Sajjad

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a giant, incredibly smart library (a Large Language Model, or LLM) that can write stories, answer questions, and solve problems. But there's a catch: the books inside are written in a secret code. You can see the words on the page, but you don't know why the librarian chose those specific words to answer your question.

For a long time, scientists tried to figure out the librarian's logic by looking at individual words (tokens). It's like trying to understand a movie by looking at single frames. You see the actors, but you miss the plot.

Other scientists tried to group these words into "concepts" (like grouping all words about "sports" or "politics") to get a better picture. But they hit a wall:

The "Super-Organizer" method (Hierarchical Clustering) was great at finding perfect groups, but it was so slow and memory-hungry that it could only handle a tiny library. If you tried to use it on a massive library, it would crash.
The "Quick-Sort" method (K-Means) was fast and could handle huge libraries, but the groups it made were often messy and didn't make much sense semantically.

Enter VQLC: The "Digital Filing Cabinet"

The authors of this paper propose a new method called VQLC (Vector Quantized Latent Concept). Think of it as building a super-efficient, digital filing cabinet for the library's secret code.

Here is how it works, using a simple analogy:

1. The Filing Cabinet (The Codebook)

Imagine you have a cabinet with 400 drawers (this is the "codebook"). Each drawer represents a specific "concept" or idea, like "Baseball," "Oil Prices," or "Toxic Comments."

2. Sorting the Mail (Vector Quantization)

When the library processes a sentence, it turns every word into a complex data point. Instead of trying to compare every single word to every other word (which is slow), VQLC takes each word and asks: "Which of my 400 drawers does this word fit into best?"

It's like a mail sorter that instantly drops a letter into the correct bin based on a zip code. This process is incredibly fast and doesn't require a massive amount of memory, no matter how many letters (words) you have.

3. Learning the Labels (Training)

At first, the drawers are empty and the labels are wrong. The system reads millions of sentences and learns to adjust the labels.

If it sees the words "Sheffield," "Ramirez," and "Jeter," it realizes they all belong in the "Baseball Players" drawer.
If it sees "GDP," "Oil," and "Earnings," it puts them in the "Market Indicators" drawer.

The system uses a "reconstruction" trick: it tries to rebuild the original sentence from just the drawer labels. If it can rebuild the sentence well, it knows it put the words in the right drawers.

4. The Result: A Clear Explanation

Once the system is trained, you can ask it: "Why did the model predict this news story is about Business?"
Instead of showing you a messy list of words, VQLC points to the specific drawers it used: "Ah, it used the 'Market Indicators' drawer because it saw words about oil prices and GDP."

Why is this a big deal?

The paper tested this new "filing cabinet" against the old methods on 12 different scenarios (using different models and datasets). Here is what they found:

Speed and Size: Unlike the "Super-Organizer" method, VQLC doesn't crash when the library gets huge. It uses very little memory, almost as little as the "Quick-Sort" method.
Quality: Unlike the messy "Quick-Sort" groups, VQLC's drawers are very clear and make sense.
The Sweet Spot: VQLC found the perfect balance. It is as fast as the cheap methods but as smart as the expensive ones.
Decoder Models: The authors noted that this method works especially well on "decoder-only" models (the kind used for chatbots like the ones we use today), where it found the clearest concepts.

The Bottom Line

The paper claims that VQLC is a scalable, efficient way to open the "black box" of AI. It turns the confusing, high-dimensional math inside a model into a neat set of understandable concepts, allowing us to see exactly what ideas the AI is using to make its decisions, without needing a supercomputer to do the math.

Note: The paper focuses strictly on interpreting how these models work (like explaining why a news classifier picked "Business"). It does not claim this method is currently used for clinical diagnosis, legal decision-making, or other specific real-world applications beyond understanding the model's internal logic.

Technical Summary: Vector Quantized Latent Concepts (VQLC)

Problem Statement

Large Language Models (LLMs) encode rich semantic information within their hidden states, yet interpreting these internal representations remains challenging. While token-level attribution methods (e.g., Integrated Gradients, SHAP) identify salient input tokens, they offer limited insight into the high-level semantic structures encoded in the model. Latent concept discovery aims to address this by grouping contextual representations into discrete, interpretable concepts. However, existing approaches face a fundamental trade-off between scalability and concept quality:

Hierarchical Clustering (e.g., LACOAT): Produces coherent concepts but suffers from quadratic memory costs ( $O(N^2)$ ), limiting its application to small datasets.
K-Means Clustering: Scales efficiently but often yields less semantically coherent concepts and provides a weaker trade-off between scalability and faithfulness.

The paper posits that a method is needed that maintains the scalability of centroid-based methods while achieving the semantic coherence of hierarchical clustering, specifically for large-scale LLMs.

Methodology: Vector Quantized Latent Concepts (VQLC)

The authors propose VQLC, a discrete concept learning framework that adapts Vector Quantization (VQ) to map continuous hidden states to a finite set of learned discrete codes. The framework operates on frozen hidden states extracted from a specific layer of an LLM.

Architecture Components:

Encoder: A lightweight residual Multi-Layer Perceptron (MLP) maps contextual token representations ( $h^{(\ell)}_{w_i}$ ) into a code space. It combines a direct linear projection with a nonlinear correction branch to preserve the original model space while enabling stable discrete assignment.
Vector Quantizer: This module maintains a learnable codebook $E = \{e_1, \dots, e_K\}$ $E = {e_{1}, \dots, e_{K}}$ . Each encoded representation is assigned to its nearest codebook vector based on squared Euclidean distance.
- Initialization: The codebook is initialized using a "farthest-first" traversal on deduplicated training representations to ensure diverse coverage of the representation space.
- Learning: The codebook is updated using an Exponentially Moving Average (EMA) of assigned encoder outputs, rather than direct gradient steps, to ensure stability. A "dead code recovery" mechanism reinitializes unused codes using encoder outputs with the largest assignment errors.
Decoder: A residual decoder reconstructs the original hidden states from the quantized vectors. The training objective minimizes reconstruction loss ( $L_{rec}$ ) and a commitment loss ( $L_{commit}$ ) to keep encoder outputs close to their assigned codebook vectors.

Concept Construction:
After training, the model is frozen. A concept-construction pass aggregates all encoded token representations assigned to the same code index $k$ to form a concept vector $v_k$ . A latent concept $c_k$ is defined as the pair $(v_k, \text{assigned tokens})$ . At inference, new token representations are assigned to the learned codebook, and the corresponding concepts explain the semantic information encoded in those representations.

Key Contributions

Scalable Discrete Learning: VQLC introduces a vector quantization framework for latent concept discovery that avoids the quadratic memory dependence of hierarchical clustering.
Optimized Trade-off: The method achieves a favorable balance between faithfulness (concept quality) and scalability, maintaining near-constant memory usage as dataset size increases.
Comprehensive Evaluation: The paper provides a multi-faceted evaluation across 12 dataset-model settings, including:
- Scalability: Runtime and memory analysis.
- Faithfulness: Measured by the change in model confidence and prediction labels after ablating the concept direction via orthogonal projection.
- LLM-as-a-Judge: Evaluation using multiple LLMs (Claude, Gemini, DeepSeek) to rank concept alignment with model predictions.
- Qualitative Analysis: Visualization of concept word clouds and case studies.
- Comparison with SAEs: Contrast with Sparse Autoencoders (SAEs) to highlight differences between feature-level and concept-level explanations.

Results

Scalability: On the ERASER dataset with the Qwen model, VQLC scales linearly with token count. While K-Means memory usage grew from 1.8GB to 17.5GB and LACOAT failed at 200k tokens (requiring ~229GB), VQLC remained nearly constant, increasing only from 1.94GB to 2.54GB at 500k tokens.
Faithfulness: VQLC achieved the highest mean confidence change in 6 out of 12 dataset-model settings. It demonstrated the clearest gains on decoder-only models (Llama-2, Qwen), outperforming baselines in 5 out of 6 decoder settings, while remaining competitive with K-Means and LACOAT on encoder-based models (RoBERTa).
LLM Evaluation: In LLM-as-a-judge rankings, VQLC achieved the best average rank in 8 out of 12 settings, suggesting its concepts are more task-aligned than those from K-Means or LACOAT.
Comparison with SAEs: Unlike SAEs, which produce a large number of sparse features with low activation rates, VQLC generates a smaller set of concepts (e.g., ~400) that are active on a significantly larger fraction of test examples. VQLC concepts were found to encode more task-relevant information for prediction in the Qwen model setting.
Qualitative Findings: Visualizations showed VQLC concepts were semantically coherent (e.g., distinct clusters for "Baseball Players" vs. "Football" in Sports, or "Clinical Drugs" vs. "Network Infrastructure" in Sci/Tech). In case studies, VQLC provided more precise explanations for model predictions compared to the broader clusters of K-Means or LACOAT.

Significance and Claims

The paper claims that VQLC offers a practical and scalable alternative to clustering-based concept discovery for LLMs. It successfully addresses the memory bottleneck of hierarchical clustering without sacrificing the semantic coherence of the discovered concepts. The authors emphasize that VQLC is particularly effective for decoder-only models, where it identifies task-relevant concept directions more effectively than existing baselines.

The work positions VQLC not as a replacement for mechanistic interpretability tools like SAEs, but as a complementary approach that organizes representations into task-level latent concepts rather than monosemantic features. The authors note limitations, including the sensitivity of the framework to hyperparameters (codebook size, commitment weight) and the current focus on classification tasks, suggesting that extending this approach to generative token-by-token reasoning is a direction for future work.