[论文解读] Analyzing Vision Transformers for Image Classification in Class Embedding Space

Martina G. Vilas, Schaumlöffel, Timothy+1 more|arXiv (Cornell University)|Oct 29, 2023

Explainable Artificial Intelligence (XAI)被引用 7

一句话总结

本文提出了一种新颖的框架，通过将中间隐藏状态投影到学习到的类别嵌入空间，来解释视觉变换器（ViTs）的机制，揭示了图像标记如何通过注意力和MLP机制逐步与类别原型对齐。该方法能够高效且准确地识别分类任务中的显著图像区域，其在捕捉与任务相关的特征方面优于线性探针。

ABSTRACT

Despite the growing use of transformer models in computer vision, a mechanistic understanding of these networks is still needed. This work introduces a method to reverse-engineer Vision Transformers trained to solve image classification tasks. Inspired by previous research in NLP, we demonstrate how the inner representations at any level of the hierarchy can be projected onto the learned class embedding space to uncover how these networks build categorical representations for their predictions. We use our framework to show how image tokens develop class-specific representations that depend on attention mechanisms and contextual information, and give insights on how self-attention and MLP layers differentially contribute to this categorical composition. We additionally demonstrate that this method (1) can be used to determine the parts of an image that would be important for detecting the class of interest, and (2) exhibits significant advantages over traditional linear probing approaches. Taken together, our results position our proposed framework as a powerful tool for mechanistic interpretability and explainability research.

研究动机与目标

• 开发一种方法，用于解释视觉变换器在图像分类过程中如何构建类别表征。
• 理解自注意力和MLP层在形成特定类别表征中的作用。
• 使用统一且高效的框架，识别对分类预测最相关的图像区域。
• 展示该框架在捕捉与模型决策相关特征方面相较于传统线性探针的优势。
• 在无需优化的情况下实现机制性可解释性，为ViT内部机制提供人类可理解的洞见。

提出的方法

• 使用输出投影矩阵，将任意层的图像标记隐藏状态投影到类别嵌入空间。
• 利用预训练的类别嵌入矩阵作为参考空间，量化内部表征与类别原型之间的对齐程度。
• 通过正样本和负样本的标记扰动实验，评估图像区域对分类的重要性。
• 通过在移除低分标记后测量模型准确率，将该框架的性能与线性探针进行比较。
• 利用自注意力中的键-值记忆对机制，分析注意力动态如何促进类别表征的形成。
• 仅通过一次前向传播在验证数据上计算可解释性得分，避免为额外分类器训练带来的高昂计算成本。

Figure 1: Schematic of our framework. (a) The hidden states of image tokens $x_{n}$ in a block $b$ are projected onto the class embedding space using the output embedding matrix $\mathbf{E}$ . (b) A key-value memory pair system in a self-attention layer. Key vectors $k_{j}$ belong to different atten

实验结果

研究问题

• RQ1 在推理过程中，视觉变换器中的图像标记如何逐步与学习到的类别原型对齐？
• RQ2 自注意力和MLP层在构建类别表征中分别发挥什么作用？
• RQ3 类别嵌入空间投影方法能否识别出对分类最具信息量的图像区域？
• RQ4 该框架在捕捉与模型决策相关特征方面与线性探针相比表现如何？
• RQ5 注意力机制和上下文信息在类别特定表征形成过程中发挥多大程度的影响？

主要发现

• 图像标记在ViT层级结构的早期阶段即发展出类别特定表征，且通过自注意力和MLP层的逐步作用，与类别原型的对齐程度不断提高。
• 该框架能有效识别对分类至关重要的图像区域，即使在移除最多60%的最不重要标记后，性能仍保持稳定。
• 正样本与负样本扰动实验表明，与随机移除相比，该方法的AUC显著更高，证实了其强大的预测能力。
• 该方法在识别与任务相关的特征方面优于线性探针，因为线性探针常会解码出可分但无关的表征。
• 该框架在时间效率上优于线性探针，仅需一次前向传播，而无需为每层进行多次训练步骤。
• 自注意力中的键-值记忆对机制被证明在类别表征形成中具有差异化贡献，注意力与MLP组件各自发挥独特作用。

Figure 2: Evolution of class identifiability mean scores across normalized blocks.