[논문 리뷰] Not all images are worth 16x16 words

이번에 소개할 논문은 제목에서 알 수 있듯이, 모든 이미지를 똑같은 patch size로 나눌 필요가 없다!라는 주제를 가지고 있다. ViT는 patch size를 달리하였을 때 더 많은 tokens이 만들어졌을 때 더 높은 성능을 보여주나 computational cost가 급격하게 증가하는 단점을 가지고 있다. 저자는 이미지가 저마다 가지고 있는 특성이 다르기 때문에 각각에 대해서 token의 개수를 dynamic하게 정해줌으로써 위와 같은 문제점을 완화할 수 있다고 주장한다. 본 논문으로 들어가보자!

 

https://arxiv.org/abs/2105.15075

Not All Images are Worth 16x16 Words: Dynamic Transformers for Efficient Image Recognition

 1. Vision Transformer(ViT)

이 논문은 ViT를 기반으로 하는 method이기 때문에 해당 논문을 먼저 읽고 오는 것을 권장한다.

 

https://aistudy9314.tistory.com/67?category=1044949 

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

 

간단하게 요약하면, 모든 픽셀들에 대해서 attention을 하는 것은 computation과 memory 측면에서 불가능에 가까울 정도로 크기 때문에 이미지를 patch 단위로 쪼개서 하나의 token 개념으로 self-attention을 적용시키는 것이 ViT이다. patch size를 작게 줄수록 더 세밀하게 볼 수 있지만 token의 개수가 많아지기 때문에 self-attention에서 사용되는 computational cost가 크게 증가하게 된다. 논문에서는 모든 이미지를 ViT처럼 16x16 patch를 적용할 필요가 없다는 것을 실험을 통해서 보여주고 있다.

 

2. Do All images are worth 16x16 words?

모든 이미지들을 똑같은 수의 patch size로 나눌 필요가 있을까? 어떻게 보면 당연하게 이미지마다 특성이 다르고 복잡성(?)(예: 흰 배경에 있는 강아지와 숲에 있는 원숭이)에서 차이가 있을 수 있기 때문에 개별적으로 다른 patch size를 주는 것이 optimal일 것이다. 

 

 

실제로 저자들이 여러 개의 token개수에 대해서 실험을 해보았는데, 14x14와 4x4의 경우 computational cost는 무려 8.5배의 차이를 보이는데 반해 accurcacy는 15.9% 정도의 차이로 상대적으로 적은 감소를 보여주고 있다. 즉, 4x4 tokens로도 충분한 "쉬운" 이미지들에 대해서 14x14 tokens을 사용하는 것은 computation 측면에서 매우 inefficient하다는 것이다.

 

3. Dynamic Vision Transformer

 

(왼쪽) overview of DVT, (오른쪽) example of DVT

 

저자들은 위와 같은 문제를 해결하기 위해서 Dynamic Vision Transformer(DVT)를 제안한다. 각 이미지에 대해서 자동적으로 몇 개의 tokens을 사용할 지 정해주는 것이 목표이며, 방법은 엄청 직관적이고 간단하다.

 

3.1 Training

다른 tokens 개수에 대해서 여러 개의 transformers를 학습시킨다는 것 외에는 기존의 방식과 똑같다(하나 학습 시키는데도 엄청 오래걸리는데...). 이는 test time에서 적은 수의 tokens에서부터 많은 수의 tokens까지 다양한 입력에 대한 networks가 필요한 것이기 때문에, 학습할 transformers의 개수는 하나의 하이퍼 파라메터가 될 것이다.

$$minimize \ \frac{1}{|D_{train}|} \sum_{(x,y) \in D_{train}}[\sum_i L_{CE}(p_i, y)]$$

 

3.2 Inference

test time에는 이미지 하나에 대해 적은 수의 tokens부터 시작해서 정해놓은 threshold를 넘길 때까지 순차적으로 token의 개수를 늘려나가는 방법을 사용한다. 위의 예제를 보면, apple의 경우 easy image로 tokens를 2x2만 사용해도 높은 confidence를 보여주고 있고 palace는 4x4 tokens를 사용해야만 91.4%의 confidence를 갖는 것을 볼 수 있다. 그리고 token의 개수는 달라지더라도 embedding dimension은 모두 똑같이 두어 더 fine-grained representation(patch들이 제각각 적절한 기능을 가지도록)을 할 수 있도록 한다.

3.3 Feature and Relationship Reuse

Training에서 다양한 token 개수에 대한 여러 Transformers를 학습시킬 때, 이전 모델에서 계산된 결과들을 그냥 버린다면 매우 inefficient할 것이다. 따라서 논문에서는 이전에 계산된 결과들을 reuse하는 feature reuse와 relationship reuse mechanisms을 제안하는데, 첫번째의 경우 이전에 학습한 deep features를 사용하고 두번째는 self-attention relationship을 이용하는 방식이다. 이는 최소한의 computational cost를 사용하여 test accuracy를 향상시키는 결과를 보여준다.

 

두 mechanisms을 설명하기 전에, 먼저 ViT의 formulation을 체크하고 넘어가자.

$$z'_l=MSA(LN(z_{l-1}))+z_{l-1}, \ \ l \in {1, \dots, L}$$

$$z_l=MLP(LN(z'_l))+z'_{l}, \ \ l \in {1, \dots, L}$$

 

MSA는 Multi-head self-attention, LN은 Layer Norm, z는 l번째 transformer layer output을 뜻한다. 위 식은 나중에 reuse mechanism이 추가된 formulation과 비교할 때 사용된다.

 

Feature Reuse

 

DVT에서 학습한 모든 Transformer는 recognition을 잘하자! 라는 공통된 목표를 가지고 있기 때문에, scratch부터 학습시키는 것보다는 이전에 계산된 deep features를 기반으로 학습을 하는 것이 더 효율적일 것이다. feature reuse mechanism은 이전 Vision Transformer의 output tokens를 이용하여 layer-wise embedding $E_l$을 학습시킨다.

$$E_l = f_l(z^{up}_L) \in \mathbb{R}^{N \times D'}, \ \ f_l: \mathbb{R}^{N \times D} \rightarrow \mathbb{R}^{N \times D} $$

 

함수 f를 더 살펴보면 nonlinearity와 flexible transformation을 위해서 LN-MLP를 사용하고, 효율성을 위해 낮은 $D'$을 준다고 한다. 그 다음으로는 reshape를 통해 기존 이미지 형태로 만들고 upsample을 한 후에 다시 token형태로 flatten을 한다. 마지막으로 나온 output을 MHA의 output과 Concat해주면 된다.

$$z_l = MLP(LN(Concat(z'_l, E_l))) + z'_l$$

**classification token은 reuse하지 않고 zero padding을 주었다고 한다.

 

Relationship Reuse

 

 

Relationship reuse는 이전 모델의 self-attention maps를 사용한다. 이전과 다르게 모든 layer의 attention logits을 Concat하여 사용한다. 이렇게 해줌으로써 각 self-attention block이 이전 모델의 모든 attention heads information을 access할 수 있게 된다.

$$ A^{up} = Concat(A_1^{up}, A_2^{up} \dots, A_L^{up}) \in \mathbb{R}^{N_{up} \times N_{up} \times N_{up}^{Att}}$$

이 방법도 layer-wise로 이루어지며 이전 attention logits을 2차원 벡터로 flatten한 후에 이를 MLP에 넣고(마찬가지로 작은 D'을 사용한다) 다시 reshape한 뒤 upsampling한다. 

$$Attention(z_l) = Softmax(A_l + r_l(A^{up}))V_l$$

여기서 $r_l$은 $A^{up}$의 information을 integrate하는 transformation network이고 위 이미지에 해당한다.

**relative importance는 logits의 크기를 조절함으로써 자동적으로 학습할 수 있다고 한다.

 

또한 일반적은 upsample operation을 적용할 수가 없기 때문에 위 이미지와 같이 upsampling을 시켜준다고 한다.

 

4. Experiments

여다른 리뷰와 마찬가지로 실험부분은 결과만 보도록 하겠다. 자세한 내용을 알고 싶다면 본 논문을 읽어보기를 ㅎㅎ..

 

DVT가 다른 models에 비해서 적은 연산량으로 높거나 비슷한 accuracy를 보여준다. 가장 성능이 높은 TNT와 같은 accurcay에서 3.6배의 연산량 차이가 나는 것을 볼 수 있다. 

 

 

위 실험은 DVT가 다른 모델들에 비해서 같은 처리량에 대해 더 높은 accuracy를 얻을 수 있다는 것을 보여준다.

 

 

reuse mechanisms을 사용했을 때와 안했을 때의 결과 차이를 보여주는 그래프이다. reuse를 아예 사용하지 않았을 때는 차이가 꽤 크고(4%), 하나의 method만 제거했을 때는 약 2%의 차이를 보여주고 있다. 

 

이외에도 더 많은 실험과 ablation study가 있지만 개인적으로 별로 중요하지 않다 생각했기 때문에 생략하였다.

---

지금까지 Not all images are worth 16x16 words 논문을 리뷰하였다. 직관적인 방법으로 computational cost를 줄인 논문이었는데, 여러 개의 ViT를 학습해야 한다는 점과 test 시에 batch를 사용할 수 없다는 점이 단점으로 다가왔다. 그리고 최악의 경우(모든 이미지가 hard samples일 때)에 모든 networks를 한 번씩 거쳐야할 수도 있다...그래도 Dynamic한 tokenization method를 처음 제안한 논문이고 결과도 뛰어나니 한번 쯤 참고할만 한 것 같다.ㅎㅎ