[논문 리뷰] Analyzing Vision Transformers for Image Classification in Class Embedding Space

Martina G. Vilas, Schaumlöffel, Timothy외 1명|arXiv (Cornell University)|2023. 10. 29.

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.