Vision Transformer (3) - Attention Map

Vision Transformer (1)

Vision Transformer (2)

---

Transformer 모델의 가장 큰 특징은 self-attention 으로 시퀀스의 각 위치가 어느 위치에 집중하는지 쉽게 시각화해서 볼 수 있다는 점입니다. 이미지 기반 딥러닝에서도 모델의 결과를 설명하려는 interpreting explanability (XAI) 시도가 매우 많고 이를 여러 포스트에서 살펴보았는데요, 이번 포스트에서는 ViT 모델의  입력 이미지에 대한 explaianability 를 attention 을 이용하여 알아보도록 하겠습니다. 또한, 실습으로 최근 페이스북에서 릴리즈한 3개의 헤드를 가진 Deit Tiny 모델을 사용하도록 하겠습니다. 사용하기 위해서는 timm 이라는 대표적인 이미지모델의 파이토치 사전훈련 모델을 제공하는 라이브러리를 먼저 설치해야 합니다. (pip install timm)

import torch

model = torch.hub.load('facebookresearch/deit:main', 
'deit_tiny_patch16_224', pretrained=True)

Attention

Attention 수행을 위해서는 Q, K, V 가 필요합니다. Deit Tiny 모델에서는 매 층마다 3개의 헤드가 존재하고 Q, K, V 는 모두 3x197x64 크기의 텐서입니다. 따라서 시퀀스 하나의 요소가 feature 차원 64를 이용하여 시퀀스 내의 모든 197개의 요소를 바라봅니다. 197개의 요소는 원본 이미지를 패치화시킨 14x14=196 과 최종 이미지의 예측을 위한 'class token' 으로 구성되어 있습니다. Attention 은 해당 모든 이미지 패치의 쿼리 $q_i$에 대해서 정보가 키 $k_j$로부터 얼마나 흐르는지를 나타내는 방법으로 Figure 1과 같이 정의됩니다.

Figure 1

Information flow

먼저 $Q\cdot K^T$ 부분은 $N\times N$=197x197 크기를 가지고 있고 이미지 패치 $i$, 채널 $c$의 쿼리 벡터 $q_{ic}$와 패치 $j$, 채널 $c$의 키 벡터 $k_{jc}$의 내적으로 이루어져 있고 softmax 함수를 거치게 됩니다. 따라서 $K$에 존재하는 패치 $j$가 가지고 있는 정보를 $Q$에 존재하는 각 이미지 패치에 얼마나 전달하는지를 정하는 것이라 생각할 수 있는데요, 내적이다보니 $q_{ic}, k_{jc}$가 같은 부호를 가지고 있다면 $A=softmax(\frac{Q\cdot K^T}{\sqrt{d_k}})$ 의 $A_{ij}$값이 커질 것이고 다른 부호를 가지고 있다면 작아질 것입니다. 즉, $A_{ij}$가 커진다면 $j$에 있는 키 정보를 $q_i$에 잘 전달할 것이고 작다면 정보를 잘 전달하지 못할 것입니다.

Figure 2는 간단하게 첫번째 헤드의 26번째 채널의 query, key 를 이미지화시켜 바라본 것인데요, ($Q, K$가 197x64 크기를 가지고 있으므로 하나의 채널을 선택한 뒤 14x14 로 reshape 하면 가능합니다.) 키 이미지는 비행기를 하이라이트하고 쿼리 이미지를 이미지 전체를 하이라이트 합니다. 쿼리 이미지의 대부분이 양수이고 키 이미지의 비행기 부분만 양수이므로 키 이미지의 비행기 부분의 정보만 쿼리 이미지 전체에 전달될 것입니다.

Figure 2

반대로 Figure 3에서는 쿼리 이미지는 비행기의 아랫 부분만 밝고 키 이미지는 비행기의 윗부분이 특히 어둡습니다. 이런 경우에는 비행기의 윗부분 정보가 대부분이 음수인 쿼리 이미지 전반에 전달될 것이고 키 이미지에서 밝은 값을 가진 비행기 이외부분의 정보는 쿼리 이미지의 비행기 아랫부분에 대해서만 전달될 것입니다.

Figure 3

Attention rollout

Layer 를 거치면서 attention map 의 양상은 residual connection 과 덧붙여져서 매번 달라질텐데 이를 ViT 처음부터 끝까지 전체 과정에 대해서 어떻게 표현할 수 있을까요? ViT 논문에서는 "attention rollout" 이라는 방법을 사용합니다. 이것은 "Quantifying Attention Flow in Transformers" 논문에서 제안된 방법으로 각 Transformer encoder 블락이 attention map $A_{ij}$ 를 가지고 있고 이것은 전 layer의 $j$번째 패치가 현 layer의 $i$번째 패치에 대해 얼마나 attention 을 부여할지 나타내므로 Transformer 전반에 거쳐 인코더 블락의 attention map 을 합치는 방법을 제안합니다.

일단 단순히 매 layer 마다 $A$를 곱하는 것을 생각할 수 있겠으나 인코더 블락에는 residual connection 이 존재하므로 시퀀스의 각 요소마다 스스로를 attention 하게끔 강제되므로 $A$에 identity 행렬을 더해준 $A+I$를 사용합니다. 그렇다면 multihead 에 대해서는 어떻게 할까요? "Attention rollout" 저자는 모든 헤드의 attention map 을 평균내어서 사용합니다. 이후 살펴보겠지만 평균 대신에 최소, 최대 등 다른 통계치를 사용해도 됩니다. 결과적으로 layer $L$에 대해서 다음 식과 같이 표현할 수 있고 특정 위치에 대한 attention 총합이 1로 만들기 위해 매번 행을 1로 만드는 정규화를 수행합니다.

$AttentionRollout_L = (A_L + I)AttentionRollout_{L-1}$

하지만 멀티헤드에 대해서 평균치를 사용하면 각 헤드별로 서로 다른 위치를 집중하기 때문에 attention map 이 전반적으로 넓게 분포합니다. 따라서 노이즈를 제거하기 위해 멀티헤드에 대해 최소값을 사용하면 Figure 4와 같이 최종 attention map 이 보다 더 국소적인 부분을 가리키고 선명해집니다.

Figure 4

여기서 더 나아가 attention score 가 낮은 부분을 일정 비율로 제거하고 멀티헤드에 대해 최대값으로만 표현하면 Figure 5와 같이 보다 선명하고 국소적인 attention map 을 얻을 수 있습니다. 특히 제거하는 비율을 높일수록 salient 객체를 나타내는 부분이 더 줄어들게 됩니다.

Figure 5

 

Implementation

개와 고양이가 같이 있는 사진에 대해서 "Deit Tiny" 모델에 대해 attention rollout 을 Pytorch 와 구글코랩을 이용해 구현해보겠습니다. 먼저 구글 드라이브를 통해 이미지를 불러오고 ViT 입력 이미지 크기에 맞게 224x224 로 바꿔준 이후에 정규화하고 'unsqueeze' 메소드를 통해 배치 축을 만들어줍니다. 이때 위에 선언한 모델에 대해서도 "cuda()" 메소드를 통해 GPU 위에 올려주어야 합니다.

import cv2
import torch

import numpy as np

from PIL import Image
from torchvision import transforms

transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
        ])
        
img = Image.open('/content/drive/MyDrive/Test/dogcat.png')
img = img.resize((224, 224))
input_tensor = transform(img).unsqueeze(0)
if torch.cuda.is_available():
    input_tensor = input_tensor.cuda()

이후에 "attention rollout" 을 구현합니다. 먼저 매 layer 마다 attention map 을 추출하기 위해 "register_forward_hook" 메소드를 이용해 모듈 이름을 이용하여 attention map 을 담습니다. "head_fusion" , "discard_ratio" 아규먼트는 각각 멀티헤드를 합치는 방법과 attention 낮은 값을 얼마나 버릴지의 비율을 정의합니다.

class VITAttentionRollout:
    def __init__(self, model, attention_layer_name='attn_drop', head_fusion="mean",
        discard_ratio=0.9):
        self.model = model
        self.head_fusion = head_fusion
        self.discard_ratio = discard_ratio
        for name, module in self.model.named_modules():
            if attention_layer_name in name:
                module.register_forward_hook(self.get_attention)

        self.attentions = []

    def get_attention(self, module, input, output):
        self.attentions.append(output.cpu())

    def __call__(self, input_tensor):
        self.attentions = []
        with torch.no_grad():
            output = self.model(input_tensor)

        return rollout(self.attentions, self.discard_ratio, self.head_fusion)

이후 "VITAttentionRollout" 클래스의 객체를 선언하고 입력 텐서에 대해 호출하면 "rollout" 함수가 실행됩니다. 이때 "self.attention" 리스트에는 12개 층에 대한 [1, 3, 197, 197] 크기의 attention map 이 담기게 됩니다. 'head_fusion' 아규먼트에서 정의한 대로 멀티헤드를 합쳐 [1, 197, 197] 크기로 attention map 을 만들고 'discard_ratio' 에서 정의한 비율만큼 Pytorch 의 'topk' 메소드를 이용하여 작은 값을 0으로 만듭니다. 이때, 'class token' 부분은 이미지 전체에 대한 representation 으로 지우지 않습니다. 매 layer $A+I$에 대해 행 정규화를 수행하고 곱해가면서 최종적인 attention map 을 출력합니다. 중요한 것은 이후 attention map 으로 마스크를 만들 때에는 이미지 전체에 대한 representation 인 'class token' 에 대한 나머지 196개 패치에 대한 attention 을 봐야한다는 점입니다. 따라서 "mask=result[0,0,1:]" 과 같이 구현되고 mask 크기는 14x14 가 됩니다.

def rollout(attentions, discard_ratio, head_fusion):
    '''
    attentions: [1,3,197,197] 크기의 텐서가 12개 담긴 리스트
    '''
    result = torch.eye(attentions[0].size(-1))
    with torch.no_grad():
        for attention in attentions:
            if head_fusion == "mean":
                attention_heads_fused = attention.mean(axis=1)
            elif head_fusion == "max":
                attention_heads_fused = attention.max(axis=1)[0]
            elif head_fusion == "min":
                attention_heads_fused = attention.min(axis=1)[0]
            else:
                raise "Attention head fusion type Not supported"

            # Drop the lowest attentions, but
            # don't drop the class token
            flat = attention_heads_fused.view(attention_heads_fused.size(0), -1)
            _, indices = flat.topk(int(flat.size(-1)*discard_ratio), dim=-1, largest=False)
            indices = indices[indices != 0]
            flat[0, indices] = 0
            
            I = torch.eye(attention_heads_fused.size(-1))
            a = (attention_heads_fused + 1.0*I)/2
            a = a / a.sum(dim=-1)

            result = torch.matmul(a, result) # [1,197,197]

    # Look at the total attention between the class token,
    # and the image patches
    mask = result[0, 0 , 1 :] # [196,196]
    # In case of 224x224 image, this brings us from 196 to 14
    width = int(mask.size(-1)**0.5)
    mask = mask.reshape(width, width).numpy()
    mask = mask / np.max(mask)
    return mask # [14,14]   

이후에는 mask 를 원래 이미지에 덧붙여 확인해보면 됩니다. 먼저 14x14 크기의 마스크를 224x224 크기로 리사이즈를 한 이후에 cv2 의 칼라맵을 적용하여 RGB 형태로 만들고 원본 이미지와 더해줍니다. 참고로 구글 코랩에서는 cv2 라이브러리의 imshow 함수를 바로 적용할 수 없으므로 코랩 자체에서 지원하는 cv2_imshow 를 사용합니다.

def show_mask_on_image(img, mask):
    img = np.float32(img) / 255
    heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
    heatmap = np.float32(heatmap) / 255
    cam = heatmap + np.float32(img)
    cam = cam / np.max(cam)
    return np.uint8(255 * cam)
    
from google.colab.patches import cv2_imshow

np_img = np.array(img)[:, :, ::-1]
mask = cv2.resize(mask, (np_img.shape[1], np_img.shape[0]))
mask = show_mask_on_image(np_img, mask)
cv2_imshow(mask)

 

참조

• https://github.com/jacobgil/vit-explain
• https://jacobgil.github.io/deeplearning/vision-transformer-explainability

 

---

Vision Transformer (4) - Pytorch 구현