Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Xuemin Yu, Ankur Garg, Samira Ebrahimi Kahou, Hassan Sajjad

TL;DR

This paper tackles the challenge of interpreting hidden representations in large language models by proposing Vector Quantized Latent Concept (VQLC), a scalable latent-concept discovery method built on the VQ-VAE framework. VQLC maps continuous token representations to a discrete codebook of latent concept vectors, enabling interpretable explanations while maintaining scalability through vector quantization, top-k sampling, and EMA-based codebook updates. It systematically compares VQLC to hierarchical clustering and K-Means, showing near-linear scalability and competitive faithfulness and interpretability across sentiment, toxicity, and news classification tasks, with extensive ablations and qualitative analyses. The work demonstrates that discrete latent concepts discovered via a learned codebook can yield finer-grained, human-understandable explanations suitable for large-scale models and real-world interpretability demands.

Abstract

Deep Learning models encode rich semantic information in their hidden representations. However, it remains challenging to understand which parts of this information models actually rely on when making predictions. A promising line of post-hoc concept-based explanation methods relies on clustering token representations. However, commonly used approaches such as hierarchical clustering are computationally infeasible for large-scale datasets, and K-Means often yields shallow or frequency-dominated clusters. We propose the vector quantized latent concept (VQLC) method, a framework built upon the vector quantized-variational autoencoder (VQ-VAE) architecture that learns a discrete codebook mapping continuous representations to concept vectors. We perform thorough evaluations and show that VQLC improves scalability while maintaining comparable quality of human-understandable explanations.

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

TL;DR

This paper tackles the challenge of interpreting hidden representations in large language models by proposing Vector Quantized Latent Concept (VQLC), a scalable latent-concept discovery method built on the VQ-VAE framework. VQLC maps continuous token representations to a discrete codebook of latent concept vectors, enabling interpretable explanations while maintaining scalability through vector quantization, top-k sampling, and EMA-based codebook updates. It systematically compares VQLC to hierarchical clustering and K-Means, showing near-linear scalability and competitive faithfulness and interpretability across sentiment, toxicity, and news classification tasks, with extensive ablations and qualitative analyses. The work demonstrates that discrete latent concepts discovered via a learned codebook can yield finer-grained, human-understandable explanations suitable for large-scale models and real-world interpretability demands.

Abstract

Deep Learning models encode rich semantic information in their hidden representations. However, it remains challenging to understand which parts of this information models actually rely on when making predictions. A promising line of post-hoc concept-based explanation methods relies on clustering token representations. However, commonly used approaches such as hierarchical clustering are computationally infeasible for large-scale datasets, and K-Means often yields shallow or frequency-dominated clusters. We propose the vector quantized latent concept (VQLC) method, a framework built upon the vector quantized-variational autoencoder (VQ-VAE) architecture that learns a discrete codebook mapping continuous representations to concept vectors. We perform thorough evaluations and show that VQLC improves scalability while maintaining comparable quality of human-understandable explanations.
Paper Structure (43 sections, 14 equations, 10 figures, 12 tables)

This paper contains 43 sections, 14 equations, 10 figures, 12 tables.

Table of Contents

  1. Introduction
  2. Problem Formulation
  3. Methodology
  4. Encoder
  5. Vector Quantizer
  6. Codebook Initialization
  7. Codebook Learning
  8. Codebook Mapping
  9. Latent Concept Content Extraction
  10. Decoder
  11. Training Objective
  12. Experiment Setup
  13. Data
  14. Models
  15. Module-specific hyperparameters
  16. ...and 28 more sections

Figures (10)

  • Figure 1: Architecture overview: Contextual token representations extracted from a deep learning model are first passed through an adaptive residual encoder to preserve the semantic structure of the input while enabling stable vector quantization. The adaptive representations are then fed into a vector quantizer, where each token is assigned to its nearest codebook vector to obtain a quantized representation. Finally, a decoder reconstructs the contextual representations from the quantized codebook vectors.
  • Figure 2: Peak Memory Usage vs. Number of Tokens
  • Figure 3: Examples of concept explanations for the RoBERTa model on the AG News dataset. Sentence denotes the input sentence, ground truth the gold label, and prediction the model output. Salient token refers to the token with the highest attributed score.
  • Figure 4: Examples of latent concepts identified in the RoBERTa model for the AG news classification task
  • Figure 5: Qualitative comparison examples of concepts on the AG News dataset using Qwen: Correct Prediction
  • ...and 5 more figures