[논문 리뷰] How do Vision Transformer Work?

이번에 소개할 논문은 ICLR 2022 spotlight를 받은 How do Vision Transformer Work? 라는 논문이다.

https://arxiv.org/abs/2202.06709

How Do Vision Transformers Work?

 

제목에서 알 수 있듯이, Vision Transformer가 왜 좋은 또는 좋지 않은 특성을 가지는 지에 대해서 분석하는 내용을 담고 있다.

 

바로 들어가보자!

---

1. Introduction

Convolutional Neural Network (CNN) vs Vision Transformer (ViT)

 

먼저 Convolutional Neural Network (CNN)과 Vision Transformer (ViT)의 특성을 간략하게만 짚고 넘어가자!

 

CNN은 local kernel을 가지고 이를 input에 sliding 방식으로 적용하여 weight sum을 통해 output feature를 만들어낸다. 또한, kernel은 모든 local input에 대해서 weight를 공유하기 때문에 translation equivalence와 같은 특성을 가진다. 

 

ViT는 먼저 이미지를 patch 단위로 나누고, 이를 tokenizing하여 모든 tokens과의 attention score를 계산한다. 

 

ViT에 대한 자세한 내용은 다음 글을 참조하길 바란다.

https://aistudy9314.tistory.com/67

[논문 리뷰] An Image is Worth 16x16 Words: Transformers for image Recognition at scale(VIT)

위와 같은 mechanism으로부터 생기는 두 architecture간의 특성 차이가 있는데,

 

첫 번째로는 data-dependency이다. CNN은 static kernel values를 가지는 반면, ViT는 data가 바뀜에 따라 attention score도 달라지게 된다.

 

다음으로는 local kernel을 통해 output이 결정되는 CNN은 local한 정보를 볼 수 있다는 특성을 가지고, 반대로 ViT는 모든 input tokens을 attention score 계산하는데 사용하므로 global하게 정보를 볼 수 있다는 특성이 있다.

 

마지막으로, CNN은 channel 단위로 kernel을 만들어주고 ViT는 input을 1차 vector로 aggregate하기 때문에, 각 channel-specific and channel-agnostic하다.

CNN	ViT
Data-agnostic	Data-specific
Local filters	Global filters
Channel-specific	Channel-agnostic

이러한 특성은 추후 실험에서 언급이 될 것이다. channel관련해서는 딱히 없긴 해요..

 

Previous Researches

Vision Transformer의 특성에 대한 기존 연구들에 대해서는 간단하게만 언급하고 넘어가도록 하겠다. 해당하는 주장들에 대한 실험들을 이 논문에서도 하기 때문에 굳이 깊게 알 필요는 없다. 

 

첫 번째는 local multi-head self attention (MSA)가 global MSA보다 더 좋은 성능을 낸다는 것이다. Long-range dependency가 자연어처리 분야에서 Transformer model의 큰 장점 중 하나였기 때문에, vision 분야에서도 이를 좋은 특성으로써 받아 들였는데, 실제로 실험을 통해 알아본 결과, self-attention을 local하게 보는 것이 성능 향상에 더 도움이 되었다고 한다.

 

두 번째는 ViTs가 data corruptions, image occlusions, 그리고 adversarial attacks에 견고한 특성을 보인다는 것이다. 그러한 이유에 대해서는 다양한 의견이 있지만, 이 논문에서는 주파수 관점에서 이야기를 하고 있다. 이는 추후 실험에서 더 알아보도록 하자.

 

마지막으로 MSA가 깊은 layer에서 사용될 수록 예측 성능에 더 큰 영향을 미친다는 것이다. CNN과 MSA를 조합하는 이전 논문들을 보면 초기 layers는 convolution이 이점을 가지고 후반 layers들은 MSA가 성능을 높이는데 큰 도움을 주는 것을 알 수 있다.

 

Questions

위에서 언급한 두 architecture의 특성과 previous researches로부터 저자들은 3가지 질문을 제기하는데, 이는 다음과 같다.

• What properties of MSAs do we need to better optimize NNs?
• Do MSAs act like Convs?
• How can we harmonize MSAs with Convs?

 

2. Question 1: What properties of MSAs do we need to improve optimization?

첫 번째로 Multi-head self-attention의 어떠한 properties가 optimization을 향상시키는 가에 대한 질문이다. 저자들은 여러 주장을 제기하면서 실험을 통해 타당성을 증명한다.

 

The stronger the inductive biases, the stronger the representations

 

Vision Transformer가 나왔던 초기에는, Transformer라는 구조가 자연어 처리에서 높은 성능을 보여주는 하나의 이유로 global dependency를 많이 언급했기 때문에, 마찬가지로 global information을 볼 수 있다는 점이 ViT가 높은 성능을 내는 이유라고 흔히 간주되었다. 하지만 저자들은 이러한 weak inductive bias (global dependency)가 학습에 이점이 별로 없다고 주장한다.

 

이를 증명하기 위해 저자들은 base ViT와 거기에 locality와 같은 constraint를 준 파생 모델 (Twins, Swin, PiT), ViT만큼 inductive bias가 약한 MLP-Mixer, 그리고 가장 Inductive bias (e.g. translation equivalance, locality)를 가지는 CNN의 test error와 training negative log-likehoold (NLL) 값을 비교하였다.

 

위 figure에서 알 수 있듯이, 오히려 ViT에 Inductive bias를 더해준 파생 모델들이 ViT보다 test error와 training NLL 모두 낮고 가장 inductive bias가 강한 CNN이 제일 좋은 성능을 보여준다. 반대로 Inductive bias가 약한 Mixer와 ViT는 상대적으로 약세를 보인다.

주의: 이 논문에서는 모든 실험에 CIFAR-100이 사용되었습니다. 이를 참고하여 실험 결과들을 보시기 바랍니다.

ViT does not overfit small training datasets.

 

Vision Transformer는 작은 데이터셋으로 학습이 매우 힘들다는 단점이 있는데, 다만 저자들은 이가 over-fitting 문제는 아니라고 말한다. 위 figure를 보면, dataset size가 줄어듦에 따라 test error와 training NLL이 같이 감소하는데 이는 over-fit이 아니라 under-fit이다.

 

ViT's non-convex losses lead to poor performance

 

저자들의 여러 실험들 중 가장 흥미로웠던 실험은 loss function의 hessian matrix eigenvalue를 통한 분석이다. 

 

저자들은 ViT가 non-convex losses을 가지기 때문에 performance가 좋지 않다고 주장하는데, 어떻게 loss가 non-convex 특성을 가지는 지 알 수 있을까? 논문에서는 hessian matrix eigenvalue를 활용하여 이를 증명한다.

 

hessian matrix란, loss function의 이차 미분 성분을 가지고 있는 matrix인데 이의 eigenvalue는 그냥 curvature (곡률: 굽어있는 정도)라고 이해하면 된다. 

 

먼저 왼쪽 figure는 warmup (맨 위)부터 300 epochs (맨 아래)까지의 eigenvalue density graph이다. 여기서 알 수 있는 점은 초기 학습 단계 (1, 2 rows)에서 ViT가 negative eigenvalues를 가진다는 것인데, 이게 바로 non-convex 특성임을 나타낸다. 반면에 ResNet은 positive eigenvalues만을 가진다.

 

딥러닝에서 loss function이 convex해야 optimization이 용이하다는 것은 많이 알려져 있는 사실이다. 즉, non-convex함을 가진다는 것은 그만큼 학습이 힘들다는 것이고 그렇기 때문에 작은 데이터셋으로 optimization이 잘 안된다고 이해할 수 있다.

 

오른쪽 figure는 큰 데이터셋을 사용하면 이러한 경향을 완화시켜준다는 것을 보여준다. Negative eigenvalues을 보면 기존 데이터셋의 6%만큼만으로 학습했을 때보다 전체 데이터셋을 사용했을 때 더 적은 negative eigenvalues를 가지는 것을 알 수 있다.

 

MSAs flatten the loss landscape

 

다음 주장은 MSAs가 loss landscape를 flat하게 만들어준다는 것이다. 이도 위와 같이 eigenvalues를 통해 알 수 있는데, 이번에는

positive eigenvalues에 집중하길 바란다. 

 

 

보면 ViT의 eigenvalues가 CNN의 eigenvalues보다 크기가 훨씬 작다. 이전의 말했듯이, eigenvalues를 curvature이고 이 값이 크다는 것은 그만큼 굽어져 있다는 것을 의미한다. 때문에 ViT가 ResNet보다 flat한 losses를 가지고 있다는 것이고, 이는 ViT가 더 높은 performance와 generlization을 갖게 해준다.

 

위 오른쪽 그림은 각 모델의 optimization trajectory를 visualize한 것인데, 위에서 말한 특성 (low curvature) 때문에 ViT가 훨씬 smooth하게 움직이는 것을 볼 수 있다.

 

A key feature of MSAs is a data specificity

이 글의 sub-section 제목은 모두 본 논문에서 가져왔는데, 개인적으로 해당 제목은 부적절하지 않나 생각이 든다. 왜냐면 여기서 하는 실험들이 data specificity를 다루는 것이 아닌 weak inductive bias (long-range dependency)가 key feature가 아니라는 것을 증명하기 때문이다. 

 

 

해당 실험에서 CIFAR-100을 사용했기 때문에 8x8 kernel이 일반적인 ViT의 patch size이고, 그 아래인 5x5와 3x3은 local하게 attention을 계산한 것이다. Figure (a)를 보면 kernel size가 오히려 작았을 때, 즉 locality라는 inductive bias를 주었을 때, error rate가 낮아지는 것을 볼 수 있다.

 

또한 이는 위의 주장과 일치하게, figure (b)를 보면, attention을 local하게 보았을 때 (5x5, 3x3)가 global하게 본 것 (8x8)보다 negative eigenvalues의 정도가 낮고, 작은 positive eigenvalues를 가진다. 

 

논문에서 "data specificity는 이전 연구에 따라서 당연하게 key feature이다" 정도로만 설명이 되어있다. 그 이유는 MLP-mixer가 input data에 dependent하지 않은데 ViT보다 성능이 좋지 않고, data에 의존성을 주고 self-attention을 사용하지 않았을 때 성능이 올라간다는 결과가 있었기 때문이라고 말한다.

(필자도 related works를 다 읽어보진 않았기 때문에 관심이 있으면 논문에서 찾아 읽어보시길....)

 

3. Question 2: Do MSAs act like Convs?

두 번째 질문은 Multi-head self-attention이 Convolutional neural network와 비슷하게 동작하는가? 이다. 물론 비슷하게 동작하면 흥미가...ㅎㅎ. 이는 다른 논문에서도 실험을 했었던 질문인데 기본적으로 이론상 MSAs는 최소 Convolution과 비슷하게 동작은 할  수 있다. 다만 inductive bias가 약하기 때문에 flexibility가 높고 그만큼 search space가 넓기 때문에 다르게 동작하는 면들이 있다.

 

MSAs are low-pass filters, but Convs are high-pass filters

 

Low-pass filters란, low-frequency만을 통과시키는 것을 의미하고 high-pass filters는 그 반대이다. 이러한 주장은 이전에도 많이 실험으로 증명 되었는데 여기서는 조금 흥미로운 실험으로 이를 이야기하고 있다.

 

왼쪽 figure는 network depth에 따른 ViT의 high-frequency에서의 (Fourier transformed) relative log amplitude 그래프이다. 어렵게 생각할 것 없이 그냥 layer의 따른 high-frequency의 정도라고 이해하면 된다 (회색은 MSAs이고 흰색은 MLPs).

 

대부분의 경우 (초기 layers를 제외), MSAs는 high-frequency 정도를 줄이고 MLPs는 반대로 높이는 역할을 하고 있다. 비슷하게, 중간 figure를 보면, 마찬가지로 ViT는 대부분의 high-frequency 정보를 낮추는 경향을 보이고, 반대로 ResNet은 높이는 것을 볼 수 있다. 

 

이러한 특성으로 인해서 high-frequency noise에는 ViT가 더 robust하고, 반면에 low-frequency noise에는 ResNet이 더 robust한 결과를 보인다 (오른쪽 figure). 

**특정 범위의 주파수만을 통과시킨다는 것은 그만큼 해당 주파수 특성을 중요하게 본다는 말과 같다.

 

MSAs aggregate feature maps, but Convs do not

 

다음으로는, MSAs가 feature maps을 aggregate하는 역할을 한다는 주장이다. 이를 보이기 위해서 feature map의 variance를 측정하였다.

 

ResNet에서 convolution은 variance를 높이는 역할을 하고 있지만, ViT에서 MSAs는 반대로 variance를 낮추는 것을 볼 수 있다. Feature map의 variance를 낮추는 것은 uncertainty를 낮추는 것이고, 이는 optimization에 도움이 되기 때문에 MSAs의 이러한 특성을 잘 살리면 학습에 도움이 될 것이다.

 

추가적으로, 두 architecture 모두 일정한 패턴을 가지는 것을 볼 수 있는데 논문에서는 이를 stage라고 명칭한다. ResNet은 각 stage의 끝 부분에서 peak가 일어나고 깊어질 수록 variance가 계속해서 커지는 데, 이를 보완하기 위해  각 stage의 끝에 MSAs를 넣어 variance를 감소시키고, 후반 layers일수록 MSAs의 head 수와 dimensions을 높게 주는 법을 생각해볼 수 있을 것이다.

 

4. Question 3: How can we harmonize MSAs with Convs?

지금까지 ViTs를 CNN과 비교하여 어떤 다른 특성을 가지는 지, 또 어떤 이점 또는 단점이 있는 지를 살펴보았다. 그럼 이러한 특성을 잘 활용해서 두 아키텍쳐를 잘 융화시킬 수 없을까?

 

 

저자들은 먼저 multi-stage NN의 layer간 correlation을 보았는데, 왼쪽 그림을 보면 특정 구간마다 block형태의 영역이 나눠져 있는 것을 볼 수 있다. 이를 통해 Neural Network가 여러 개의 개별적인 모델들을 사용하는 것처럼 multi-stage로 구성되어있으며, 이 stage 단위로 디자인을 적용하는 것이 합리적이라고 생각할 수 있다.

**E: embedding, P: Pooling

 

두 번째로 lesion study (해당 layer의 특정 노드를 제거한 후 결과를 체크)를 진행하였는데, ResNet의 경우 깊은 layers를 제거하는 것보다 초반 layers를 제거했을 때 성능 감소가 심하고, 반대로 Swin은 초반 layers에선 MLPs를 제거했을 때, 후반 layers에선 MSAs를 제거했을 때 성능 저하가 크게 일어난다. 즉, Conv나 MLPs는 초반 layers에서 큰 역할을 하고, MSAs는 후반 layers에서 성능에 더 큰 영향을 미친다는 것을 알 수 있다.

 

Blind-up rule

위를 토대로, 저자들은 blind-up rule을 만들어 이를 바탕으로 새로운 hybrid-architecture를 제안하며 alterNet이라고 명칭하였다.

1. Alternatively replace Conv blocks with MSA blocks from the end of a baseline CNN model.
2. If the added MSA block does not improve predictive performance, replace a Conv block located at the end of an earlier stage with an MSA block.
3. Use more heads and higher hidden dimensions for MSA blocks in late stages.

https://github.com/xxxnell/how-do-vits-work

 

논문에는 해당 rules만 써져있는데, 솔직히 한번에 그 과정을 이해하기가 좀 까다롭다 ㅎㅎ. 찾아보니 github에 그 과정을 잘 그린 figure가 있어 가져왔으니 참고하기 바란다. 

 

간단하게만 설명하면, ResNet50 (또는 다른 CNN 모델)을 base로하고 후반 layers부터 하나 씩 Convolution을 MSA로 대체한다. 만약 성능이 오르지 않았을 경우, 이전 (shallower) stage의 끝에서부터 다시 MSA으로 Convolution을 대체한다. 

 

 

간단한 방법이지만 CIFAR-100 (small dataset)에서 SOTA 성능을 보여준다. 더불어 위에서 주장했던 것과 일치하게 robustness도 높고, Hessian matrix eigenvalue 관점에서도 baseline인 ResNet보다 작은 positive eigenvalues를 가진다.

**Negative eigenvalues가 생기는 것은 문제가 되지 않을까 생각을 하였는데, 성능이 좋다하니 뭐...작은 negative eigenvalues만을 가져서 큰 영향을 미치지 않는다거나?

 

---

Vision Transformer가 도대체 왜 성능이 좋은 지, 왜 optimization이 힘든지를 다양한 실험을 통해 분석하여 좀 더 ViT를 이해하는데 도움을 준 좋은 논문이었다. 인공지능은 성능 싸움 또는 아이디어 싸움을 많이 하는데, 이렇게 왜? 라는 여러 의문점을 해결해주는 분석 논문들도 많이 나와주면 좋을 것 같다.

 

아쉬운 점은 위에서도 언급했지만 CIFAR-100에서만 실험이 되었다는 것이다. ImageNet에서도 똑같은 결과가 나올지는 장담을 못하기 때문에 신뢰성이 약간은 떨어지지만 그럼에도 여러 intuition을 주는 좋은 논문임에는 틀림없다.