[Paper Review] ICLR 2024, Interpreting CLIP's Image Representation via Text-Based Decomposition 논문 리뷰

ICLR 2024, Interpreting CLIP's Image Representation
via Text-Based Decomposition

논문 링크: https://openreview.net/attachment?id=5Ca9sSzuDp&name=pdf
프로젝트 페이지 링크: https://yossigandelsman.github.io/clip_decomposition/

---

 

1. Introduction

최근 많은 논문들에서 거대한 Text, Image Pair로 학습시킨 CLIP 모델을 활용하는 후속 연구들을 진행하고 있다.
본 논문에서는 CLIP의 이미지 인코더를 분석하여 모델의 각 구성 요소가 final representation에 미치는 영향을 해석 가능한 텍스트로 설명하고 있다.

먼저, Attention Head의 역할을 파악하기 위한 "TEXT SPAN"이라는 알고리즘을 고안하여,
각 Attention Head가 위치(location)나 형태(shape) 중 어떤 속성에 집중하고 있는지를 자동으로 찾아냈다. 
Figure 1의 (a)에서 볼 수 있듯이 같은 이미지에 대해 layer 22의 head 1에서
"A semicircular arch", "An isosceles triangle", 그리고 "An oval"과 같은 것들이 
top 3 basis에 포함되어 있는 것으로 보아, 해당 attention head는 shape에 집중하고 있음을 알 수 있다.
"TEXT SPAN"으로 얻은 text basis를 활용하면 text를 통해 이미지의 spatial structure를 얻을 수도 있게 된다.
예를 들어 Figure 1의 (c)에서와 같이 피라미드의 삼각형 부분에 "isosceles traiangle"이라는 label이 높은 연관성을 보이는 것으로 보았을 때,
spatial decomposition이 이미지의 어떤 영역에 영향을 주고 있는지를 확인할 수 있다. 

또한, CLIP에서 부적절한 특징(Spurious Feature)를 제거하여 낮은 성능을 내고 있는 데이터셋에 대해 성능이 올라감을 확인했다.
이와 더불어, CLIP 내에서 Spatial Localization 역할을 수행할 수 있음을 밝혔고,
이를 활용하여 Figuare 1의 (b)에서와 같이 Zero-shot Image Segmentation을 수행할 수 있는 방법을 제시했다.

이렇게 CLIP의 요소들을 text-interpretable하게 분석하는 과정을 통해 Transformer 모델의 이해 가능성을 확장할 수 있으며,
이를 통해 모델을 수정하고 개선할 수 있음을 보여주고 있다.

 

 

---

3. Decomposing CLIP Image Representation into Layers

해당 파트에서는 CLIP 모델의 Image Representation을 계산하는 방법과
이를 개별 레이어로 decomposition하는 과정에 대해서 설명하고 있다.

 

3.1 CLIP ViT Preliminaries

Contrastive Pre-training.

CLIP은 Text의 representation과 Image의 representation 간의 similarity를 비교하여
i = j인 경우에는 similarity score가 높아지고 i != j인 경우에는 score가 낮아지도록 학습된다.

Zero-shot Classification.

대규모 Image, Text Pair로 학습된 Text Encoder와 Image Encoder를 활용하여 Zero-shot Prediction을 수행할 수 있다.
"A photo of a {object}"라는 템플릿을 활용하여 Text Encoder에 입력으로 넣어주게 되며,
Image Representation 과의 similarity score가 가장 높은 것을 선별하는 방식으로 classifiation task를 수행할 수 있다.

CLIP Image Representation.

Vision Transformer (ViT)는 Image를 d-dimensional representation으로 projection한다.
CLIP의 Image Encoder로 사용되는 ViT에서 Image Representation $M_{image}(I)$는
vision-and-language의 joint space로 linear projection을 한 $d^{'}$-dimensional representation을 뜻한다.
아래 수식에서 P는 projection matrix를 뜻하여 training 시에 학습되는 matrix이다.

ViT Architecture.

ViT는 L개의 layer로 이루어져 있으며, 각 레이어는 multi-head self-attention(MSA) + MLP 로 이루어져 있다.
이미지 I는 N개의 non-overlapping image patch들로 나누어지며,
각 patch들은N개의 d-dimensional vector로 linear projection된 후  position embedding과 합쳐진다.
class token $z^{0}_{0}$은 additional learned token으로 추후에 output token으로 사용된다.

token $z^{0}_{0}, z^{0}_{1}, z^{0}_{2}, ...$의 column으로 이루어진 matrix $Z^{0} \in R^{d \times (N + 1)}$은
아래와 같이 L iteration 동안 업데이트 되는 두 가지 residual step으로 이루어져 있다.
residual stream $Z^{l}$의 첫 번째 column은 class token $[Z^{l}]_{cls}$과 연관되어 있으며,
따라서 ViT의 output은 $[Z^{L}]_{cls}$가 된다.

 

3.2 Decomposition into Layers

ViT의 Residual Structure는 모델의 output을 각 layer의 direct contribution에 대한 합산으로 표현할 수 있도록 한다.
Image Representation $M_{image}(I)$가 ViT output의 linear projection임을 고려할 때,
아래와 같이 Image Representation은 MLP와 MSA 레이어의 기여도(direct contributions)로 분해될 수 있다.
여기서 한 레이어에서 다른 downstream layer로 이어지는 output에 대한 indirect contributions는 무시했다고 한다.
최종적으로, CLIP의 representation을 분석하기 위한 decomposition을 수행할 때 아래와 같은 식이 사용되었다고 한다.

Evaluating the Direct Contribution of Layers.

위와 같은 식을 통해서 final representation을 결정하는 direct effect들의 대부분은 late attention layer에서 온다는 것을 찾아냈다.
각 direct effect들의 역할을 알아내기 위해 각 component들을 데이터셋에 포함된 이미지들의 평균값으로 대체하는 mean-ablation 실험을 진행했다고 한다.
특히나 ablation을 수행하기 전과 후의 zero-shot classification accuracy의 변화를 집중적으로 관찰하며,
어떤 component가 없을 때 가장 큰 성능 하락이 일어나는지 확인했다.
실험에서 사용된 데이터는 ImageNet validation set의 classification 성능을 확인했으며,
사용된 모델은 LAION-2B로 학습된 OpenCLIP ViT-H-14, L-14, 그리고 B-16을 사용했다.

MLPs have a Negligible Direct Effect.

MLP 레이어들에 대한 mean-ablating 결과를 살펴보았을 때,
MLP는 image representation에 직접적인 영향은 거의 미치지 않는다는 것을 확인할 수 있었다.
실험 결과 1~3% 내외의 작은 수치 하락만 존재했다.

Only the Last MSAs have a Significant Direct Effect.

MSA 레이어들에 대해서 mean-ablating을 수행했다.
초반 MSA 레이어들을 제거했을 때는 큰 성능 하락이 있지 않았지만,
마지막 MSA 레이어들을 제거했을 때는 성능이 크게 떨어짐을 확인했다.
이를 통해서 마지막 4개의 MSA 레이어가 최종 output에 큰 영향을 미치고 있다는 것을 알 수 있었다.
따라서 추후 분석을 할 때, MLP나 early MSA 레이어들에 대한 분석보다는
마지막 4개 MSA 레이어들을 집중 분석하는 것을 택했다.

 

3.3 Fine-grained Decomposition into Heads and Positions

MSA 블록들에 대해서 fin-grained decomposition을 수행했다.
위에서 설명한 decomposition 식에서와 같이 class token의 output 결과에 집중하여 분석했다.
해당 식에서 MSA의 결괏값은 H 개의 independent attention head와 N 개의 input token들 간의 합산으로 치환하여 생각할 수 있다.
여기서 $W^{l, h}_{V O} \in R^{d \times d}$는 transition matrices를 뜻하며,
$\alpha^{l, h}_{i} \in R$은 class token에서 얻은 attention weight를 뜻한다.
따라서 MSA ouput은 각각의 individual head와 token들의 direct effect로 표현될 수 있음을 뜻한다.

이를 이전에 설명한 decomposition 식에 대입하면 아래와 같은 식을 얻을 수 있다.

이를 달리 말하면, 모든 attention block의 direct effect들은 모든 dimension에 걸친 tensor c의 합이라는 것을 알 수 있다.
만약, 몇몇 dimension에 대해서만 확인한다면, effect들을 다양한 방법으로 decompose할 수 있게 된다.
예를 들어서 spatial dimension $i$에 대해서만 contract하여 각 head에서의 contribution을 구할 수 있다.
$$c^{l, h}_{head} = \Sigma^{N}_{i=0} c_{i,l,h}$$
반대로, layer와 head 방향으로 contract하게 되면 각 image token들에 대한 contribution을 구할 수 있다.
$$c^{i}_{token} = \Sigma^{L}_{l=1} \Sigma^{H}_{h=1} c_{i,l,h}$$

이때 $c^{l,h}_{head}$나 $c^{i}_{token}$에 해당하는 $c_{i,l,h}$는 joint text-image representation space 내에 있으므로, 우리가 이를 text로 해석할 수 있게끔 되는 것이다.
예를 들어서 주어진 text description t에 대해서 $$<M_{text}(t), c^{l,h}_{head}>$$를 계산한다면 각 head에서 text description t와의 similairty를 직관적으로 확인할 수 있게 된다.