[Review] TokenLearner: What Can 8 Learned Tokens Do for Images and Videos

요약: 기존의 ViT, ViViT를 활용할 때 있어서, 패치의 높은 연산량 문제를  TokenLearner module이라는 효율적으로 해결할 수 있는 토큰 학습 방법을 제안하였음

이 논문을 이해하기 위해서는 아래 3가지 동향을 알고 있어야 한다.

기존에 텍스트 시퀀스나 시계열적 데이터를 처리했던 Transformer모델을 이미지에 적용하는 방법을 제안한 것이 ViT고, 비디오에 적용하는 방법을 제안한 것이 ViViT이다. 논문에서는 이러한 Transformer 기반의 모델을 비전 분야에 적용할 경우 만날 수 있는 문제를 다음과 같이 정의하면서 시작했다.

Introduction

1) Main Challenge in Vision Transformers - Patch calculation quantity

image patch

논문에서 정의한 문제는 다음을 따른다.

1. Vision Transformer 구조는 종종 합리적인 결과를 도출하기 위해서 너무 많은 토큰을 필요로 한다는 것
2. 만약 512X512 이미지를 16x16 토큰으로 분할하려면 총 1024개의 토큰이 생성됨

이거를 비디오에 적용한 ViViT의 경우엔 어떻게 될까? 비디오는 사진을 1초에 여러 장을 빠르게 보여주는 매체이므로 이미지와 마찬가지로 토큰으로 분할할 경우 quadratic하게 계산량이 증가한다.

Video

2) Proposed method preview

• 위의 문제들을 해결하기 위해 저자는 Token을 고정된 형태로 나누어 사용하는 것보다 적응적으로(adaptively) 적은 수의 토큰을 생성하는 일종의 학습 모듈을 제안
• 생성된 토큰들은 Transformer 아키텍처 내 자유롭게 배치가 가능하며 후속 레이어에서 계속 쓰면서 토큰을 줄여 나갈 수 있음

요약: TokenLearner를 통해 얻을 수 있는 이점

1. 입력 값에 대해 동적으로 token을 생성 가능 $\to$ 이미지의 표현 벡터로도 볼 수 있음
2. 비디오의 경우 토큰 수를 효과적으로 줄임으로써 연산량을 줄일 수 있음
3. 모듈은 transformer encoder안에 자유롭게 위치시킬 수 있음

 

Method

1) Notation

$X_t$: 시공간 형상의 입력 텐서, t=1이면 image, t>1이면 video

$Z_t=[z_i]_{i=1}^S$: 입력 텐서는 S개의 토큰으로 분할될 수 있음, 각 토큰은 $z_i$로 표기하며 입력 텐서의 차원이 HxWxC라면 토큰은 차원이 1x1xC이다. 또한 이는 Tokenizer function $A_i(X_t)$로 투영 됨($z_t=A_i(X_t)$)

2) Token Learner Modules - tokenizer function $A_i(X_t)$

$A_i(\cdot )$를 학습하여 $X_t$의 픽셀의 유익한 조합을 적응적(adaptively)으로 선택할 수 있다. 이는 사전에 설정한 토큰 개수 S개 만큼 존재한다. 생성된 토큰은 입력텐서의 시공간적 관계, 상호작용을 의미하며, 입력 텐서의 고정 분할이 아니라 적응적으로 변화하는 선택적인 공간 집합이 된다. 즉, 다시 말해 입력텐서를 고정된 크기로 동일하게 모두 쪼개는 것이 아니라 마치 하이퍼파라미터 처럼 사람이 지정해서 몇개로 쪼갤 수 있는 것이다.

이는 기존의 ViViT에서 사용했던 Uniformly frame sampling 기법이나 Tubelet embedding 기법과 차이가 있다.

• $A_i(X_t)$의 구성요소

1.  Spatial attention mechanism($\alpha_i(\cdot)$): 입력 텐서를 기반으로 weight map을 계산
2.  Broadcasting function($\bigodot\gamma (\cdot ) $): element-wise multiplication
3.  Spatial global average pooling($\rho(\cdot ) $): 차원의 수를 줄이는 역할

$$z_i=A_i(X_t)=\rho (X_t\bigodot \gamma (\alpha _i(X_t)))$$

1. Spatial attention mechanism($\alpha_i(\cdot)$)

컨볼루션 계열에 의해 구현되고, sigmoid function을 사용한다. 저자는 이를 MLP layer로도 나중에 업데이트 했다고 밝히고 있다. 

2. Broadcasting function($\bigodot\gamma (\cdot ) $) & Spatial global average pooling($\rho(\cdot ) $)

생성된 weight map에 입력 텐서를 element wise 한 후 차원을 줄이는 단계이다. 출력물은 토큰으로, S개의 Cx1x1 벡터이다. 비디오의 경우 T개의 frame이 존재하므로 전체 ST개의 토큰이 생성된다. 이 모든 과정을 토큰이 학습된다는 의미에서 "TokenLearner"라 명칭했다고 밝히고 있다. 생성된 토큰을 바탕으로 ViViT와 마찬가지로 transformer encoder에 넣고 Task를 수행한다.

2) Token Fuser

Token Fuser 모듈은 Transformer layer를 통과한 후 사용하는 후속 레이어며, 선택적으로 배치할 수 있다.

저렇게 토큰 별로 선형 변환을 하는 것이 [1]에서 모델의 메모리를 크게 절약할 수 있는 이점을 얻을 수 있다고 밝히고 있고, 본 논문에서 이를 적용했다고 볼 수 있다. CNN의 경우 채널 별로 서로 다른 필터를 적용하는데, 이 대신에 채널을 모두 같은 파라미터를 공유하는 weight를 적용하여 파라미터의 수를 크게 줄일 수 있다. 이러한 방식으로 모델을 구성한 사례는 "isotropic Resnet", "pyramidal ViT"가 있다고 한다. 결론적으로 Token Fuser의 결과물은 원래 인풋 텐서와 같은 사이즈의 텐서로 복원이 된다. 이는 궁극적으로 TokenLearner module을 계속적으로 이어 붙일 수 있게 된다.

Expermiment

TokenLearner module의 효과를 확인하기 위해 이미지 학습능력을 평가했다. 두가지 버전의 모델을 가지고 실험을 진행했다.

(a)  transformer model에 TokenLearner를 삽입 → image representation task, video representation task

(b) TokenLearner 뿐 아니라, Token fuser를 transformer 내에 여러 위치에 삽입 → video representation task

1) Experiments with images

테이블에서 볼 수 있듯이, tokenlearner를 사용할 경우 계산량과 이미지 분류 정확도에서 기존 ViT계열의 모델보다 성능 향상을 이루었다.

2) Experiments with Videos

이미지 표현 task와 마찬가지로, 비디오 표현 능력을 검증하기 위해 ViViT 파생 모델을 Backbone으로 두고, TokenLearner를 삽입했을 때 성능을 비교했다. 테이블을 통해 알 수 있듯, TokenLearner를 삽입하고 layer를 더 추가했음에도 정확도, 연산량 측면에서 우수했다.

비디오 데이터셋 Kinetics-600에 대한 결과 테이블이다. 제안 모델이 정확도, 연산량 측면에서 우수했으나 tokenFuser를 사용하게 되면 연산량이 매우 증가하는 것을 알 수 있다.

Conclusion

본 논문이 가장 크게 기여한 부분은?

✓ ViT의 토큰 임베딩 방법을 video에서 쓸 수 없고, ViViT의 토큰 임베딩 방법을 image에서 쓸 수 없는 것과 반대로, 본 논문에서는 하나의 임베딩 방법으로 차원만 신경 써주면 image, video에 둘 다 사용이 가능한 TokenLearner를 제안한 것이 효율성 측면에서 우수하다고 생각함
✓ 이미지를 토큰화 하는 과정에서 FC layer를 사용하게 되면 파라미터의 수가 증가하는데(ViT), GAP를 사용한 부분에서 연산량을 줄이려고 하는 노력이 많이 보였음

본 논문의 문제점 혹은 부족한 부분은?

✓ 실험을 통한 인사이트 제시가 없음. 실험을 통해 알 수 있듯이, TokenLearner는 다른 모델과 비교 했을 때, 정확도, 연산량 측면에서 모두 우수 했지만, Fuser를 함께 쓰게 되면 비록 연산량은 크게 증가했지만 정확도는 뛰어나게 상승하였음. 
✓ 궁극적으로 Fuser가 왜 이렇게 정확도를 높여주는지 저자의 생각이 기술되지 않았기 때문에 그 이유가 궁금함

참고문헌

[1] Tolstikhin, Ilya O., et al. "Mlp-mixer: An all-mlp architecture for vision." Advances in Neural Information Processing Systems 34 (2021): 24261-24272.

[2] Vaswani, Ashish, et al. "Attention is all you need." Advances in neural information processing systems 30 (2017).

[3] Dosovitskiy, Alexey, et al. "An image is worth 16x16 words: Transformers for image recognition at scale." arXiv preprint arXiv:2010.11929 (2020).

[4] Arnab, Anurag, et al. "Vivit: A video vision transformer." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2021.