CORTEX: Concept-Oriented Token Explanation in Vector-Quantized Generative Model

I would like to remind the authors that the rebuttal phase is coming to a close, and as of now, reviewers have yet to recieve any replies to the published reviews.

If the authors look to partake in any meaningful discussion on their work, responding to the reviews are necessary. For my own sake, I would appreciate some time digesting their responses to be able to make a final decision on their work.

[W5] The methodology is somewhat unclear. Highlighting the importance of input dependence/independence and how the two approaches interact could improve clarity.

Response: The saliency-based method provides fine-grained, instance-specific explanations by identifying tokens crucial for a concept in a particular image. In contrast, the optimization-based method offers a global perspective by finding token combinations representing a concept across all possible images, independent of specific inputs.

Under circumstances where an explanation is needed for a specific generated image, the saliency-based method is more informative. When seeking a general understanding of how a concept is represented in the model's codebook, the optimization-based method is more appropriate.

Therefore, the saliency-based method interprets tokens at the level of specific generated images, while the optimization-based method explains the entire codebook. Additionally, the results of the optimization-based method can also be used for tasks such as image editing.

[W6] The method is specific to VQGMs, and diffusion models have shown higher quality in image generation, potentially limiting the impact of the work.

Response: While diffusion models are indeed powerful, VQGMs offer advantages such as efficient discrete representations and are still widely used in practice [1, 2, 3, 4, 5, 6, 7]. Furthermore, our methods can potentially be adapted to interpret discrete components within diffusion models or other generative models that incorporate quantization.

We will discuss this possibility in the revised paper, highlighting how our approach can contribute to the broader field of generative modeling, including models that combine discrete and continuous representations.

[Q3] Can mined concepts from the codebook be applied in other applications with discrete token-based backbones?

Answer: Yes, our approach is broadly applicable to any model that uses discrete tokens as part of its architecture. The core idea is to leverage a small external model (information extractor) to interpret high-dimensional token representations that are otherwise difficult for humans to understand. This makes the framework highly adaptable to other domains.

For example, in natural language processing (NLP), models with discrete vocabularies can benefit from similar analysis. By training an information extractor tailored to specific concepts in the NLP domain, token-level contributions to particular outputs can be effectively analyzed and interpreted. This demonstrates the versatility of our framework and its potential applications beyond the context of generative models.

[1] Esser, Patrick, Robin Rombach, and Bjorn Ommer. "Taming transformers for high-resolution image synthesis." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021.

[2] Wang, Xinlong, et al. "Emu3: Next-token prediction is all you need." arXiv preprint arXiv:2409.18869 (2024).

[3] Yu, Lijun, et al. "Magvit: Masked generative video transformer." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[4] Yu, Lijun, et al. "Language Model Beats Diffusion--Tokenizer is Key to Visual Generation." arXiv preprint arXiv:2310.05737 (2023).

[5] Kondratyuk, Dan, et al. "Videopoet: A large language model for zero-shot video generation." arXiv preprint arXiv:2312.14125 (2023).

[6] Jin, Yang, et al. "Unified language-vision pretraining with dynamic discrete visual tokenization." arXiv preprint arXiv:2309.04669 (2023).

[7] Jin, Yang, et al. "Video-lavit: Unified video-language pre-training with decoupled visual-motional tokenization." arXiv preprint arXiv:2402.03161 (2024).

[W1, Q4]: The saliency method based on Simonyan et al. (2013) has dubious interpretive value due to issues highlighted by Adebayo et al. (2018).

Response: In our paper, visualization is not used as an evaluation metric. Instead, we focus on quantitative comparisons, as shown in the experiments in Table 2, where we demonstrate the effectiveness of our approach by masking tokens with high Token Saliency Values (TSV) and comparing them to masking tokens based on frequency.

This experiment clearly shows that our information extractor can identify tokens strongly correlated with the target concept.

Moreover, it is important to note that the information extractor is not the model we aim to interpret. Our experiments are not designed to prove that the information extractor perfectly encapsulates all token information. The information extractor is adaptable and can be tailored to different concepts based on user selection, serving as a flexible tool for explanation.

The role of TSV in our work is to validate the effectiveness of our approach. Specifically, it shows that a small external model (the information extractor) can interpret high-dimensional, complex token information that is otherwise difficult for humans to understand. Our paper's key contribution lies in demonstrating the effectiveness of this method rather than proving the comprehensiveness of the information extractor itself.

[W2, Q2]: The use of a synthetic dataset may limit generalizability. Could the reliance on this dataset be dropped?

Response: Our method does not rely exclusively on synthetic datasets and can be applied to real-world datasets.

The primary reason we used a synthetic dataset generated by VQGM is to analyze the behavior of the VQGM. By focusing on model-generated data, we can better understand how the VQGM encodes and expresses these concepts. This emphasis aligns with our goal of explaining the VQGM’s behavior, which is why we ultimately chose to analyze data generated by the model.

Additionally, in our preliminary experiments, we tested the approach using real datasets. By encoding existing images through a pretrained encoder and using the vector quantization step to identify the closest codebook token for each embedding, we can generate the inputs required for the information extractor described in our paper. However, we ultimately chose to use generated data because using real datasets explains the distribution of the concept in real images rather than the model's understanding of the concept, which is the focus of our study.

[W3]: The value of the proposed method seems niche and specific, with limited context and baselines.

Response: We would like to clarify that the proposed method is broadly applicable and not limited to a specific niche. Currently, generative models primarily fall into two categories: diffusion-based models and vector-quantized generative models (VQGMs). Our approach is designed to work with any VQGM that employs a discrete codebook, regardless of the number of concepts or the size of the codebook.

This interpretability framework is equally applicable to text-to-image VQGMs, where the concepts can be defined as the context provided in the model's prompt. By training a lightweight information extractor tailored to the specific concept of interest, token-based concept explanations can be generated.

To the best of our knowledge, our method is the first to interpret tokens in the codebook. As such, our method is versatile, capable of interpreting any model that utilizes a codebook, and provides a generalizable approach for explaining the learned representations of any concept. This broad applicability makes the method relevant to a wide range of applications beyond the specific use case presented in our work.

[W4, Q1]: The fundamental motivation behind the work is insufficiently established. The importance of interpretability in VQGMs could be more strongly emphasized.

Response: We appreciate this feedback. Our work enhances the interpretability and transparency of generative models, aiding in debugging, bias detection, and ethical use. By revealing how concepts are encoded at the token level, developers can improve model design, control the generative process, and ensure intended behavior. Our methods address challenges such as mitigating biases, improving robustness to adversarial attacks, and enabling content customization. By providing insights into concept representation, our approach promotes responsible use, enhances transparency, and fosters user trust.

[W1] The comparison between CNN-based and ResNet-based extractors may not provide valuable insights into the effectiveness of the proposed approach.

Response: Our approach to training the information extractor is designed to ensure it learns token-level features, essentially functioning as a "knowledge storage unit." It is expected that different extractor architectures will capture different patterns, which is entirely natural. However, our goal is not to explain the information extractor itself but rather to leverage the knowledge it encodes to provide human-understandable interpretations of high-dimensional tokens.

Importantly, the extractor is not fixed; it is trained specifically for the concepts that the user is interested in. For example, if the concept involves positional information, the extractor architecture should be sensitive to spatial features. Our experiments are not intended to demonstrate that a CNN-based extractor can capture all token information. Instead, they aim to validate the effectiveness of our method: that even with a simple extractor architecture, we can interpret the tokens in VQGM, which are otherwise difficult for humans to understand.

Our proposed method and experiments are designed to demonstrate that our extractor can learn token-based information and transform it into human-understandable explanations. The goal is not to prove that our extractor has learned all possible information. The choice of information extractor depends on the specific concept being targeted, as different concepts may require different extractors.

[W2] The usefulness of the explanations obtained is unclear. Clarification on identifying biases, improving robustness, etc., is needed to justify the design.

Response: We want to clarify that our explanations are useful in applications including image editing, bias detection, etc. -As for the optimization-based method, as shown in Figure 4 of the paper, the explanations generated can be directly applied to image editing. By leveraging the results of the optimization process, specific aspects of an image can be controlled, altered, or refined, showcasing the practical utility of this approach. -The saliency-based method provides explanations for the token inputs of already generated images. In cases where the concept being interpreted is undesirable, such as bias (e.g., related to skin tone or race), the identified tokens can be edited, by deleting or modifying the corresponding tokens in the codebook, and the VQGM can be re-trained to remove these biases.

[W3] Experiments are limited to ImageNet with class-conditioned generation. The text-to-image setup is not explored, where explanation might be more important.

Response: We want to clarify that our proposed explanation method is applicable to any VQGM with a discrete codebook, regardless of the number of concepts or the size of the codebook. We demonstrate that with a small external model (the information extractor), we can generate human-understandable interpretations for VQGM tokens. This interpretability approach can be applied to any VQGM model that utilizes a codebook. Moreover, our concepts can be defined as the context for the model's prompt in a text-to-image model. By training a simple information extractor tailored to the concept that needs to be interpreted, token-based concept explanations can be obtained. Therefore, our method is broadly applicable to any model utilizing a codebook and can provide interpretations for any concept.

[Q1] Can visualization of the token space be provided for better presentation of the explanation process?

Answer: Thank you for your suggestion. The reason we did not use token space visualization is that our explanation focuses on making tokens understandable from a human visual perspective, such as identifying which tokens correspond to specific features. Our method aims to explain the relationship between tokens and concepts rather than clustering or analyzing the relationships within the high-dimensional token space. Therefore, our approach emphasizes providing explanations that are visually interpretable for humans rather than visualizing the token space itself.

[Q2] Have the authors tried masking the tokens before decoding to pixels?

Answer: Yes, we have considered masking the tokens before decoding. Removing a token would corrupt its corresponding spatial region in the decoded image, replacing it with visual artifacts. This introduces additional uncertainty and does not differ significantly from simply masking the corresponding token region in the decoded image. Since each token corresponds to a fixed region in the image, directly masking the decoded region provides the same effect while being more controlled and consistent. Therefore, we chose to mask the token regions in the decoded images with black to achieve a more interpretable and reasonable comparison.

[Q3] What would be a reasonable setup to evaluate the proposed methods on text-to-image models?

Answer: As mentioned in [W3], our method is applicable to any VQGM that utilizes a codebook, including text-to-image models. Following the approach described in our paper, the text prompt can be treated as the concept to be explained. By training a corresponding information extractor for this concept, the extractor can identify the token combinations associated with the given concept. Thus, our framework can be seamlessly applied to any text-to-image model and can provide explanations for any concept of interest.

[Q4]: How to present empirical and quantitative results to justify the usefulness of the explanations? What new capabilities could be enabled?

Answer: As mentioned in [W2], our saliency-based method can identify tokens in already generated images that correspond to concepts related to biases. By editing these specific tokens in the codebook, we can mitigate or eliminate such biases. Additionally, our optimization-based method enables image editing by optimizing token combinations to achieve desired modifications. These capabilities demonstrate the practical usefulness of our explanations, enabling both bias mitigation and fine-grained control over generated images.

[W1] The saliency-based method may risk overfitting to specific instances, reducing utility for broader applications.

Response: While the saliency-based method focuses on individual images, we mitigate overfitting by aggregating results across multiple images to identify common patterns and tokens that are consistently important for a concept. This aggregation helps generalize the explanations beyond specific instances.

Moreover, our optimization-based method complements this by providing a global, input-independent perspective, identifying tokens that are generally important for a concept across the entire codebook. Together, these methods offer both specific and general insights, enhancing the overall utility of our approach.

[Q1] Can you provide more details on how CORTEX compares to existing explainability methods for generative models?

Answer: Current generative models primarily fall into two categories: the mainstream diffusion-based models and another widely adopted approach that uses pre-trained codebooks combined with transformers for next-token prediction to generate images or videos [1, 2, 3, 4, 5, 6, 7]. The latter, which leverages vector quantization and next-token prediction from a codebook, offers advantages such as faster generation and stronger continuity between tokens, making it one of the reasons for the growing popularity of VQGMs.

Our method, CORTEX, is specifically designed for generative models utilizing codebooks, such as VQGMs. For generative models that do not use codebooks, like Stable Diffusion, our approach is not applicable. Importantly, CORTEX is the first method to propose interpreting the tokens within the codebook of a VQGM. In contrast, existing explainability methods [8, 9, 10] typically focus on pixel-level explanations rather than the token-level interpretability we aim to achieve. This distinction highlights the novelty of our approach and its specific applicability to token-based generative models.

[Q2] How do you address potential biases in the pretrained VQGMs affecting the explanations? What steps can be taken to ensure fairness and unbiased explanations?

Answer: We can address biases in pretrained VQGMs by treating bias-related concepts, such as skin tone or race, as the target concepts for explanation within our framework. By using our methods, we can identify the tokens associated with these bias-related concepts and uncover unintended correlations arising from biased training data.

Our methods can help mitigate these biases in several ways:

1. Identifying and Adjusting Tokens - By revealing which tokens are linked to biased concepts, we can adjust the codebook or modify the token selection process during generation. - For instance, if certain tokens disproportionately influence the generation of sensitive concepts, we can modify, reweight, or remove these tokens to reduce their impact.
2. Rebalancing Training Data - Retraining or fine-tuning the VQGM with more balanced and representative data can help address underlying biases in the model.

[1] Esser, Patrick, Robin Rombach, and Bjorn Ommer. "Taming transformers for high-resolution image synthesis." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021.

[2] Wang, Xinlong, et al. "Emu3: Next-token prediction is all you need." arXiv preprint arXiv:2409.18869 (2024).

[3] Yu, Lijun, et al. "Magvit: Masked generative video transformer." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[4] Yu, Lijun, et al. "Language Model Beats Diffusion--Tokenizer is Key to Visual Generation." arXiv preprint arXiv:2310.05737 (2023).

[5] Kondratyuk, Dan, et al. "Videopoet: A large language model for zero-shot video generation." arXiv preprint arXiv:2312.14125 (2023).

[6] Jin, Yang, et al. "Unified language-vision pretraining with dynamic discrete visual tokenization." arXiv preprint arXiv:2309.04669 (2023).

[7] Jin, Yang, et al. "Video-lavit: Unified video-language pre-training with decoupled visual-motional tokenization." arXiv preprint arXiv:2402.03161 (2024).

[8] Simonyan, Karen. "Deep inside convolutional networks: Visualising image classification models and saliency maps." arXiv preprint arXiv:1312.6034 (2013).

[9] Selvaraju, Ramprasaath R., et al. "Grad-cam: Visual explanations from deep networks via gradient-based localization." Proceedings of the IEEE international conference on computer vision. 2017.

[10] Nguyen, Anh, Jason Yosinski, and Jeff Clune. "Multifaceted feature visualization: Uncovering the different types of features learned by each neuron in deep neural networks." arXiv preprint arXiv:1602.03616 (2016).

[Q1] Could the values in various dimensions of the embedding affect the TSV results? What if all code entries are normalized into unit vectors?

Answer: The TSV values are calculated based on the gradient of the output with respect to each input token, and these tokens are directly drawn from the codebook. During pretraining, the tokens in the codebook are optimized based on Euclidean distance. The main goal of our study is to interpret the pre-trained tokens in VQGM codebook. Normalizing the embeddings into unit vectors would alter the original information encoded in the tokens, making it impossible to reconstruct the original image using the VQGM decoder faithfully. This is why we cannot use unit vectors to calculate TSV.

[Q2] Will the authors consider other extractor architectures like transformers?

Answer: Our approach to training the information extractor is designed to ensure it learns token-level features, essentially functioning as a "knowledge storage unit." It is expected that different extractor architectures will capture different patterns, which is entirely natural. However, our goal is not to explain the information extractor itself but rather to leverage the knowledge it encodes to provide human-understandable interpretations of high-dimensional tokens.

Importantly, the extractor is not fixed; it is trained specifically for the concepts that the user is interested in. Our experiments are not intended to demonstrate that a CNN-based extractor can capture all token information. Instead, they aim to validate the effectiveness of our method: that even with a simple extractor architecture, we can interpret the tokens in VQGM.

[Q3] What happens if we choose two semantically distant concepts or remove the mask constraint in the optimization-based method?

Answer: Thank you for your insightful suggestion. The reason we focus on semantically similar concepts in our experiments is to demonstrate that our information extractor can learn fine-grained, token-level distinctions. For example, when optimizing for one bird species, we show that modifying a specific region can lead the pre-trained model to classify the image as another bird species. This highlights that our model has effectively captured the unique characteristics of each bird species, further validating the feasibility of our method.

As for removing the mask constraint, it is certainly possible. However, in this case, as illustrated in Figure 4 with our image editing example, the entire image will turn into noise at the beginning of optimization, making it unrelated to the original input. The purpose of this experiment is to show that by optimizing only a specific region, we can increase the pre-trained model's prediction confidence for the target label. This demonstrates that our information extractor successfully captures fine-grained differences between concepts.

[Q4] Is there related analysis on the different sensitivities of ViT and CNN models?

Answer: As mentioned in Q2, our experiments are not designed to prove that the information extractor can learn all token information, nor are they aimed at interpreting the information extractor itself. Instead, we demonstrate that even a simple CNN-based information extractor can effectively capture concept-related information stored in tokens.

Additionally, our information extractor is not fixed; it is trained specifically for the concept of interest to the user. Therefore, the primary goal of our experiments is to validate that the proposed information extractor can learn meaningful information from the tokens and that our interpretability approach is feasible and effective.

[Q5] How can these methods improve the controllability of VQGMs?

Answer: By identifying tokens critical for specific concepts, our methods provide different utilities in the context of VQGMs. The optimization-based method can be directly applied to image editing by modifying token combinations to control the presence or absence of certain features in generated images. For example, as demonstrated in Figure 4, the optimization-based method allows targeted edits, such as changing the species of a bird or altering the style of an object, enabling fine-grained customization and enhancing controllability in the generation process.

In contrast, the saliency-based method identifies tokens within existing images that are associated with specific concepts. This can be used for tasks like detecting and addressing biases in the model. For instance, if tokens related to certain biases are identified, they can be removed from the codebook, and the model can be re-trained to mitigate those biases, improving fairness and reducing unwanted artifacts in the generated outputs.

[W1] Lack of context about the source and performance of the VQGAN used.

Response: We used the pretrained VQGAN model [1], trained on ImageNet. This model achieves competitive performance in image generation tasks, effectively capturing complex visual features. The performance of the VQGAN is a well-recognized conditional image generation model, but it is out of the scope of our paper. Therefore, we did not add it to our discussion. The detailed performance report of VQGAN can be found in [1].

[W2] The domain is relatively narrow. Further discussions on generalization, effects of varying concepts and codebook size, and applications to other tasks like text-to-image would be beneficial.

Response: First, we want to clarify that our approach is applicable to any VQGM with a discrete codebook, regardless of the number of concepts or the size of the codebook.

Regarding your concern about the domain being relatively narrow, it’s worth noting that current generative models primarily fall into two categories: diffusion or vector-quantized transformers. While diffusion models denoise random noise iteratively, VQ-transformers encode images into discrete tokens for autoregressive generation, offering faster inference and better semantic coherence [1, 2, 3, 4, 5, 6, 7].

This interpretability approach can be applied to any VQGM model that utilizes a codebook. Moreover, our concepts can be defined as the context for the model's prompt in a text-to-image model. By training a simple information extractor tailored to the concept that needs to be interpreted, token-based concept explanations can be obtained. Therefore, our method is broadly applicable to any model utilizing a codebook and can provide interpretations for any concept.

To the best of our knowledge, our method is the first to interpret tokens in the codebook. We demonstrate that with a small external model (the information extractor), we can generate human-understandable interpretations for VQGM tokens.

[W3] Further analysis of the relationship between saliency-based and optimization-based methods is needed, rather than discussing them side by side.

Response: Thank you for this suggestion. The saliency-based method provides fine-grained, instance-specific explanations by identifying tokens crucial for a concept in a particular image. In contrast, the optimization-based method offers a global perspective by finding token combinations representing a concept across all possible images, independent of specific inputs.

Under circumstances where an explanation is needed for a specific generated image, the saliency-based method is more informative. When seeking a general understanding of how a concept is represented in the model's codebook, the optimization-based method is more appropriate.

Therefore, the saliency-based method interprets tokens at the level of specific generated images, while the optimization-based method explains the entire codebook. Additionally, the results of the optimization-based method can also be used for tasks such as image editing.

We will revise our paper and add this discussion to the final version.

[1] Esser, Patrick, Robin Rombach, and Bjorn Ommer. "Taming transformers for high-resolution image synthesis." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021.

[2] Wang, Xinlong, et al. "Emu3: Next-token prediction is all you need." arXiv preprint arXiv:2409.18869 (2024).

[3] Yu, Lijun, et al. "Magvit: Masked generative video transformer." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[4] Yu, Lijun, et al. "Language Model Beats Diffusion--Tokenizer is Key to Visual Generation." arXiv preprint arXiv:2310.05737 (2023).

[5] Kondratyuk, Dan, et al. "Videopoet: A large language model for zero-shot video generation." arXiv preprint arXiv:2312.14125 (2023).

[6] Jin, Yang, et al. "Unified language-vision pretraining with dynamic discrete visual tokenization." arXiv preprint arXiv:2309.04669 (2023).

[7] Jin, Yang, et al. "Video-lavit: Unified video-language pre-training with decoupled visual-motional tokenization." arXiv preprint arXiv:2402.03161 (2024).