💬 NLP

Cross-Layer Discrete Concept Discovery for Interpreting Language Models

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

Published 2026-06-11

📖 5 min read🧠 Deep dive

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

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 a large language model (like the AI behind this conversation) as a massive, multi-story factory. Raw materials (words) enter at the bottom floor, get processed, and move up to the top floor where the final product (an answer or a decision) is created.

For a long time, scientists trying to understand how this factory works have only looked at one single floor at a time. They'd peek into a room on the 3rd floor and say, "Ah, this machine is doing something!" But they missed a crucial detail: the factory has a special conveyor belt system (called the "residual stream") that carries items from the bottom floor all the way to the top, mixing them with new items along the way.

Because of this conveyor belt, if you only look at one floor, you see a messy jumble of duplicate and mixed-up parts. You can't tell which part came from where or what it actually means.

The Problem: The "Blurry Photo"

Previous methods tried to clean up this mess, but they treated the information like a smooth, continuous liquid (like water). They tried to find patterns in the water, but because the water is so fluid, the patterns kept splitting apart and blending together. It was like trying to identify specific ingredients in a smoothie just by looking at the liquid; you know it's there, but you can't separate the strawberry from the banana easily.

The Solution: The "CLVQ-VAE" Factory Upgrade

The authors of this paper built a new tool called CLVQ-VAE. Think of this as a smart sorting machine that sits between two floors of the factory.

Here is how it works, using a simple analogy:

The Adaptive Encoder (The Smart Conveyor): Instead of just grabbing the raw materials from the lower floor, this machine gently adjusts them. It's like a chef who takes a raw ingredient and gives it a tiny, precise tweak to make it ready for the next step, without changing its fundamental nature.
The Discrete Bottleneck (The Stamping Machine): This is the most important part. Instead of letting the information flow as a blurry liquid, this machine forces the information to be "stamped" onto one of a finite set of pre-defined stamps (called a codebook).
- Imagine you have a box of 1,000 specific LEGO bricks.
- When the machine sees a complex idea, it doesn't try to build a new shape out of clay. Instead, it says, "This idea is most like Brick #42."
- This turns a messy, continuous blur into a clear, discrete label. It forces the AI to say, "I am using the 'Anger' brick," or "I am using the 'Sports' brick," rather than a vague mix of both.
The Decoder (The Reconstructor): The machine then tries to rebuild the "top floor" of the factory using only these specific LEGO bricks. If it can successfully rebuild the top floor using only the "Anger" brick, then we know for sure that "Anger" was a critical concept for the AI's decision.

Why This Is Better (The Results)

The paper tested this new machine against old methods (like just looking at one floor or using the "liquid" approach) on three different tasks: movie reviews, detecting toxic comments, and sorting news articles.

Here is what they found:

The "Removal Test" (Faithfulness): When the researchers took the specific "LEGO bricks" (concepts) the machine found and removed them from the AI's brain, the AI's performance crashed. In some cases, the AI got 93% worse at its job. This proves the machine found the actual reasons the AI was making decisions, not just random noise.
The "Judge Test" (LLM Evaluation): They asked other AI models to act as judges and look at the concepts found by the new machine versus the old ones. The new machine's concepts were ranked first 66.7% of the time. The judges felt these concepts made more sense and were more coherent.
The "Human Test" (Interpretability): They showed word clouds (visualizations of the concepts) to human volunteers.
- When shown the new machine's concepts, humans could guess what the AI was thinking 78% of the time.
- When shown the old method's concepts, humans only guessed correctly 54% of the time.
- Furthermore, different humans agreed much more with each other when looking at the new machine's results, meaning the concepts were clearer and less confusing.

The Secret Sauce

The paper highlights a few "tricks" that made this work:

Temperature Sampling: Instead of always picking the single "best" brick, the machine sometimes picks from the top 5 closest bricks. This is like giving the machine a little bit of randomness to explore different options, which prevents it from getting stuck using the same few bricks over and over.
Spherical vs. Flat: They found that organizing the bricks based on the direction they point (spherical) rather than just their distance (flat) helped humans understand the concepts better, even if the "flat" method was slightly better at predicting the AI's raw numbers.

Summary

In short, this paper introduces a way to translate the messy, overlapping thoughts of a large AI into a clean, organized list of distinct "concepts" (like LEGO bricks). By doing this across multiple layers of the AI, they can clearly see what the AI is thinking and why it makes its decisions, making the "black box" of AI much more transparent and understandable for humans.

Technical Summary: Cross-Layer Discrete Concept Discovery for Interpreting Language Models

Problem Statement

Interpreting large language models (LLMs) is hindered by the nature of transformer residual streams, which linearly mix and duplicate features across adjacent layers. Single-layer analysis methods often fail to capture this cross-layer structure, missing computational patterns that are only visible when multiple layers are examined together. While recent cross-layer approaches like Sparse Autoencoders (SAEs) address layer mixing, they operate in continuous spaces. In these spaces, concepts "split" across many neurons without clear boundaries, forcing arbitrary thresholds or multi-vector combinations to isolate activations. This misalignment with human reasoning, which relies on discrete categories, limits the interpretability of such methods.

Methodology

The authors propose CLVQ-VAE (Cross-Layer Vector Quantized-Variational Autoencoder), a framework designed to discover discrete, interpretable concepts by mapping representations from a lower layer ( $l$ ) to a higher layer ( $h$ ) through a discrete vector-quantization bottleneck. The framework consists of three core components:

Adaptive Residual Encoder: To preserve semantic information while enabling codebook reassignment, this component applies a learnable interpolation mechanism to input embeddings from layer $l$ . It uses a mixing coefficient $\alpha$ (constrained to $[0, 0.5]$ ) to balance the original embedding with a linearly transformed version, preventing the destruction of linguistic features while allowing targeted refinements for cross-layer reconstruction.
Vector Quantizer: This acts as a discrete bottleneck, mapping continuous encoder outputs to a single codebook vector from a learned set $E = \{e_j\}_{j=1}^K$ $E = {e_{j}}_{j = 1}^{K}$ . To ensure stable training and effective codebook utilization, the authors employ:
- Initialization: K-means (Euclidean or Spherical) or random initialization.
- Stochastic Sampling: Instead of deterministic selection, the model samples from the top- $k$ nearest codebook vectors using a temperature-controlled probability distribution. This encourages uniform codebook utilization and concept diversity.
- EMA Updates: Codebook vectors are updated using Exponential Moving Averages (EMA) of assignment counts and accumulated vectors rather than direct backpropagation, ensuring stability.
Transformer Decoder: A bidirectional transformer decoder reconstructs the target activations from layer $h$ using the quantized codebook vectors. It utilizes cross-attention mechanisms to leverage information from the unquantized encoder outputs, offloading low-level reconstruction details so the codebook can focus on distinct, high-level concepts.

The model is trained via a combined loss function: a reconstruction loss (Mean Squared Error between the decoder output and the target higher-layer representation) and a commitment loss (encouraging encoder outputs to commit to codebook vectors).

Key Contributions

Discrete Cross-Layer Mapping: Unlike standard VQ-VAEs that reconstruct the same layer, CLVQ-VAE acts as a transcoder, collapsing redundant residual-stream features from lower layers into compact, interpretable concept vectors in higher layers.
Stochastic Sampling Mechanism: The introduction of top- $k$ temperature-based sampling combined with EMA codebook updates provides controlled exploration of the discrete latent space, maintaining codebook diversity and preventing collapse.
Comprehensive Evaluation: The framework is evaluated across encoder-based models (RoBERTa, BERT) and decoder-only models (LLaMA-2-7B, Qwen2.5-3B-Instruct) on three datasets (ERASER-Movie, Jigsaw, AGNews).

Results

The paper evaluates CLVQ-VAE against clustering, single-layer VQ-VAE, and SAE baselines across three axes:

Faithfulness (Concept Ablation): Removing identified concepts from sentence representations significantly degrades model accuracy. CLVQ-VAE achieved the lowest perturbed accuracy in 5 of 12 settings, with concept removal causing accuracy drops of up to 93.2% (on RoBERTa/ERASER-Movie). It consistently outperformed SAEs, which showed inconsistent or negligible drops, particularly on decoder models.
Semantic Coherence (LLM-as-a-Judge): An ensemble of LLMs ranked CLVQ-VAE concepts first in 66.7% of comparisons, with a mean rating of 1.890, outperforming single-layer VQ-VAE, SAE, and clustering.
Human Interpretability: In a human study, annotators could recover model predictions from CLVQ-VAE visualizations with 78.2% accuracy, compared to 54.1% for clustering. Annotators also reported higher confidence (8.44 vs. 5.98) and achieved "almost perfect" inter-annotator agreement ( $\kappa = 0.864$ ) with CLVQ-VAE visualizations.

Additionally, the authors found a trade-off in initialization strategies: K-means initialization tended to yield higher functional faithfulness (better perturbation results), while Spherical K-means produced concepts with higher semantic coherence and human alignment.

Significance

The paper claims that CLVQ-VAE provides a robust framework for translating opaque model mechanisms into faithful, interpretable concepts. By integrating discrete representation learning with cross-layer analysis, the approach addresses the limitations of continuous SAEs (concept splitting) and single-layer analyses (missing cross-layer structure). The results validate that the discovered concepts are not only functionally critical to model predictions but also align with human understanding, offering a pathway to more transparent and accountable language models. The authors note that while the method is effective, it requires substantial memory for activation extraction and separate training for different layer combinations.