[論文レビュー] 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層が、クラス固有の表現を形成する上で果たす役割を理解すること。
• 統一的で効率的なフレームワークを用いて、分類に最も関連する画像領域を同定すること。
• 線形プローブに比べて、タスク関連の特徴をより効果的に捉えることができるかを実証すること。
• 最適化を伴わないで、人間が解釈可能な洞察を提供するメカニズム的解釈を可能にすること。

提案手法

• 出力投影行列を用いて、任意の層における画像トークンの隠れ状態をクラス埋め込み空間に投影する。
• 事前学習済みのクラス埋め込み行列を基準空間として用い、内部表現がクラスプロトタイプとどの程度一致しているかを定量化する。
• 正例および負例のトークン摂動実験を実施し、分類に寄与する画像領域の重要性を評価する。
• 低スコアのトークンを削除した後のモデル精度を測定することで、本フレームワークと線形プローブの性能を比較する。
• 自己注意におけるキー・バリュー記憶ペア機構を用いて、注目ダイナミクスがカテゴリカル表現形成にどのように寄与するかを分析する。
• 追加の分類器のトレーニングを避けるために、検証データに対して1回のフォワードパスのみで解釈スコアを計算する。

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がランダム削除に比べて顕著に高く、予測能力の高さが確認された。
• 線形プローブに比べて、本手法はタスク関連の特徴をより効果的に同定でき、線形プローブは分離可能ではあるが関係のない表現を解釈する傾向がある。
• 本フレームワークは線形プローブよりも時間的に効率的であり、1回のフォワードパスで十分で、各層ごとに複数回のトレーニングステップを要しない。
• 自己注意におけるキー・バリュー記憶ペア機構は、カテゴリカル表現形成に差別的に寄与しており、注目とMLPコンponentがそれぞれ異なる役割を果たしていることが示された。

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