🤖 machine learning

Concept-Centric Token Interpretation for Vector-Quantized Generative Models

This paper introduces CORTEX, a novel framework that enhances the interpretability of Vector-Quantized Generative Models by identifying concept-specific token combinations through both sample-level and codebook-level analysis, thereby improving transparency and enabling applications like targeted image editing.

Original authors: Tianze Yang, Yucheng Shi, Mengnan Du, Xuansheng Wu, Qiaoyu Tan, Jin Sun, Ninghao Liu

Published 2026-04-07

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Tianze Yang, Yucheng Shi, Mengnan Du, Xuansheng Wu, Qiaoyu Tan, Jin Sun, Ninghao Liu

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 magical art machine (like a high-tech version of a paint-by-numbers kit) that creates beautiful pictures. This machine doesn't paint with brushes; it builds images out of tiny, invisible Lego bricks called tokens.

The problem is, we don't really know which bricks make up a "cat" or a "doctor." We just know the machine spits out a picture, but the internal recipe is a black box. Sometimes, this machine gets things wrong or shows bias (like always painting white doctors and rarely painting black doctors), and we can't tell why because we can't see the recipe.

This paper introduces CORTEX, a new tool that acts like a "magnifying glass" for these invisible Lego bricks. It helps us understand exactly which bricks the machine uses to build specific ideas.

Here is how CORTEX works, explained with some fun analogies:

1. The Problem: The "Black Box" Cookbook

Think of the machine's Codebook as a giant cookbook with 10,000 pages. Each page is a specific visual "ingredient" (a token).

When you ask the machine to draw a "dog," it picks a handful of these pages and combines them.
The Issue: If you just look at the final picture, you can't tell which pages were the dog's ears, which were the background sky, and which were just random noise.
The Bias: If you ask for a "doctor," the machine might secretly rely on "white doctor" pages 4 times more often than "black doctor" pages, even if you didn't specify the race. We need a way to count these pages to prove the bias exists.

2. The Solution: The "Reverse Engineer" (Information Extractor)

To solve this, the authors built a special helper robot called the Information Extractor.

Normal Machine: You give it a word (e.g., "Bird"), and it builds a picture.
CORTEX's Robot: You give it a picture (or a pile of Lego bricks), and it has to guess, "Is this a bird? Is it a cat?"
Why this helps: To get really good at guessing, the robot has to learn exactly which Lego bricks are the most important for identifying a bird. It learns to ignore the background sky and focus on the beak and wings. This is based on a principle called the Information Bottleneck—it forces the robot to throw away the junk and keep only the essential clues.

3. The Two Superpowers of CORTEX

CORTEX uses this robot to do two different jobs:

A. The "Spot the Difference" Detective (Sample-Level)

Imagine you have a specific photo of a dog. CORTEX asks the robot: "If I cover up this specific Lego brick, will you still know it's a dog?"

If covering up a brick makes the robot say, "I don't know what this is anymore," that brick is critical.
CORTEX highlights these critical bricks in red.
Real-world use: It found that when the machine draws a "doctor," it uses "white doctor" bricks much more often than "black doctor" bricks, even when you just ask for a generic doctor. This proves the machine has a hidden bias.

B. The "Master Chef" (Codebook-Level)

Instead of looking at one picture, CORTEX looks at the entire cookbook (the Codebook).

It asks: "If I want to build the perfect 'Indigo Bird' from scratch, which specific pages of the cookbook do I need to combine?"
It mathematically searches through all 10,000 pages to find the perfect combination of bricks that creates a bird, ignoring everything else.
Real-world use: This allows for precise editing. You can tell the machine, "Keep the body of the bird, but swap the head bricks for a different species." The machine can then change just the head without messing up the rest of the image.

4. Why This Matters

Before CORTEX, these AI art machines were like magic boxes: you put a prompt in, and a picture came out, but you didn't know why it looked that way.

Transparency: CORTEX opens the box and shows us the recipe.
Fairness: It helps us catch biases (like the doctor example) so we can fix them.
Control: It lets us edit images surgically by swapping specific "ingredients" rather than just regenerating the whole picture.

In short: CORTEX is the translator that teaches us the secret language of AI art machines, turning a mysterious black box into a transparent, controllable, and fair tool.

1. Problem Statement

Vector-Quantized Generative Models (VQGMs) (e.g., VQGAN, DALL-E) have become state-of-the-art tools for high-quality image generation by utilizing a discrete latent space represented by a codebook of tokens. However, the internal mechanics of these models remain a "black box." Specifically:

Lack of Interpretability: It is unclear which specific tokens in the codebook are critical for generating a specific visual concept (e.g., "doctor," "bird," "red crown").
Bias and Shortcut Learning: VQGMs often exhibit systematic biases (e.g., generating white doctors more frequently than black doctors even with neutral prompts) or rely on shortcut features.
Limitations of Existing Methods: Traditional token selection methods often rely on frequency counts, which fail to distinguish between essential concept tokens and irrelevant background/contextual tokens, leading to cluttered and inaccurate explanations.

The core challenge is to identify concept-relevant token combinations while filtering out background noise to understand how VQGMs encode and generate specific visual concepts.

2. Methodology: CORTEX Framework

The authors propose CORTEX (Concept-Oriented Token Explanation), a framework grounded in the Information Bottleneck (IB) principle. The goal is to compress input data (tokens) while retaining the maximum information relevant to a specific semantic label.

The framework consists of three main components:

A. Information Extractor Module (IEM)

Function: Unlike generative models that map concepts to images, the IEM works in reverse, mapping token-based embeddings back to semantic concepts.
Training: The IEM is trained as a supervised classifier. It takes token embeddings ( $E$ ) generated by the VQGM and predicts the probability distribution over a set of concepts ( $Y$ ).
Role: By forcing the model to predict concepts from tokens, the IEM learns to identify and extract the most informative features (tokens) while discarding irrelevant background information, adhering to the Information Bottleneck principle.

B. Sample-Level Explanation Method

This method analyzes individual generated images to determine token importance.

Token Importance Score (TIS): For a specific image and concept, the method computes saliency scores (gradients) with respect to the IEM's prediction.
Selection: It identifies the top- $n$ tokens with the highest TIS for that specific image.
Aggregation: By aggregating these high-TIS tokens across multiple images of the same concept, it identifies the most frequent and consistent tokens ( $T^*_{concept}$ ) that define the concept.

C. Codebook-Level Explanation Method

This method explores the entire codebook space to find optimal token combinations without relying on pre-existing generated images.

Optimization: It formulates token selection as an optimization problem. A token selection matrix $P$ is optimized to maximize the IEM's prediction probability for a target concept.
Differentiability: Since token selection is discrete, the authors employ Gumbel-Softmax to make the selection process differentiable, allowing for gradient-based optimization.
Masking: A mask ( $M_{mask}$ ) allows for localized optimization (e.g., editing only the head of a bird), enabling the discovery of fundamental token combinations that characterize a concept.

3. Key Contributions

Novel Framework (CORTEX): The first approach to interpret VQGMs by identifying concept-specific token combinations using an Information Extractor based on the Information Bottleneck principle.
Dual-Level Explanation:
- Sample-Level: Provides image-specific insights into which tokens drive a specific generation.
- Codebook-Level: Discovers global, fundamental token combinations that define a concept across the entire codebook.
Bias Detection & Quantification: Demonstrated the ability to quantitatively detect demographic biases (e.g., race and gender in professional roles) by analyzing the frequency of concept-specific tokens in neutral prompts.
Targeted Image Editing: Enabled precise, localized image editing by optimizing token selections in specific regions of the codebook embedding.

4. Experimental Results

The authors evaluated CORTEX on various pre-trained VQGMs (VQGAN) using the ImageNet dataset (1,000 classes).

Validation of Token Relevance:
- Counterfactual Masking: When tokens identified by CORTEX were masked in generated images, pre-trained classifiers (ResNet, ViT) showed a significant drop in accuracy and probability for the target concept.
- Comparison: CORTEX outperformed random token selection and frequency-based baselines. For example, masking CORTEX-selected tokens reduced ResNet50 accuracy by ~46%, whereas random masking only reduced it by ~42%.
Bias Detection:
- In experiments with the prompt "A doctor in the hospital," the model generated images where tokens associated with "White doctor" appeared ~3.7 times more frequently (8.55 vs. 2.32 tokens per image) than those for "Black doctor."
- Similar disparities were found for gender (Male vs. Female doctors), confirming systematic biases in the model's latent space.
Image Editing:
- The codebook-level method successfully transformed specific features (e.g., changing a bird species by modifying tokens in the head region) while preserving the rest of the image, demonstrating precise controllability.

5. Significance and Impact

Transparency: CORTEX bridges the gap between discrete latent representations and human-understandable concepts, making VQGMs more interpretable.
Bias Mitigation: It provides a quantitative tool to detect and measure demographic biases in generative models, which is a critical step toward developing fairer AI systems.
Control: The ability to identify and manipulate specific tokens enables targeted image editing without retraining the model, offering practical utility for creative and forensic applications.
Generalizability: The framework is applicable to various VQGMs and can be extended to other discrete latent space models, including video and vision-language models.

In conclusion, CORTEX offers a robust, mathematically grounded approach to "opening the black box" of vector-quantized generative models, revealing how they encode concepts, where they fail (bias), and how they can be precisely controlled.