https://arxiv.org/pdf/2210.09461
DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification
https://arxiv.org/pdf/2106.02034 최근에 Token Pruning에 대해서 볼일이 있어서 다시 보는 김에 리뷰를 남겨봅니다. AbstractVision Transformer에서 가장 정보량이 많은 일부 토큰만 사용해서, 정확한 이미지 인식
doonby.tistory.com
- 논문을 읽기전에, 간단하게 보고 오면 좋을 token pruning 논문리뷰
Abstract
학습없이 token pruning만큼 빠르면서도, 높은 정확도를 유지할 수 있는 token merging (ToMe)를 소개합니다.
이미지에서는 throughput을 약 2배, 비디오에서는 약 2.2배까지 올릴 수 있었으며, 정확도 하락은 0.2~0.3%를 유지했습니다.
심지어 ToMe는 inference뿐만 아니라, 학습에도 적용할 수 있습니다.
다양한 실험을 통해서, 이미지, 비디오, 오디오 분야에서 높은 정확도와 속도를 달성할 수 있음을 보였습니다.
Introduction
이전에는 token pruning으로 입력 토큰수를 줄여서 속도를 향상시키는 방법들이 연구됐습니다.
하지만 학습중에는 zero-masking(실제로 토큰을 줄이지 않고, 1과 0으로 만드는)을 사용했기 때문에, 학습에는 속도향상을 볼 수 없었습니다. 또한, 토큰 정보를 지우는 것이기 때문에 정보 손실이 크게 발생하여 줄일 수 있는 토큰수에 한계가 존재합니다.
이런 큰 문제점들을 개선하기 위해, 논문의 저자들은 토큰을 제거하지 않고 합치는 Token Merging (ToMe)를 제안합니다.
ToMe는 학습과 추론 모두에서 사용할 수 있으며, 별도의 학습없이도 적용할 수 있다는 장점이 있습니다.
논문의 큰 contribution은 다음과 같습니다.
- 학습과 상관없이 ViT의 throughput과 실제 학습속도를 향상시키는 방법을 제안.
- 이미지뿐만 아니라, 동영상, 오디오에도 적용할 수 있음을 보임.
- 시각화를 통해서, 이미지와 비디오 등에서 병합 기준을 보임.
Token Merging
Token pruning처럼 Token merging도 모듈을 추가합니다.
Token Merging - Strategy
토큰을 줄이는 기준은 매번 $r$ 개의 토큰이 감소하도록 설계했습니다.
예를 들어, 총 $L$개의 모듈에서 병합을 수행하면 토큰은 $rL$개 줄어들게 됩니다. (개수입니다. 비율이 아니예요.)
장점은 token pruning에서 입력에 따라 dynamic하게 토큰 수를 줄이던 연구에서 정확도는 높지만, 배치 추론이 어려운 문제를 해결 할 수 있습니다. (동일하게 $rL$개가 줄었으니까요)
또 다른 점은, token pruning은 블록의 시작부분에 모듈을 삽입하여 토큰을 줄이는 방법을 사용했지만, ToMe는 위 Figure1처럼 attention과 MLP 사이에 삽입합니다. 이는 attention의 결과를 바로 가져와서, 토큰간의 정보를 효과적으로 사용하여 병합하기 위함이라고 합니다.
Token Merging - Token Similarity
토큰을 병합하려면 '비슷하다'는 기준을 먼저 정의해야합니다.
단순하게 토큰의 feature vector간의 거리로 정의할 수 있지만, 저자들은 이것이 최적이 아닐 수 있다고 이야기합니다.
(transformer는 overparameterized되어있어서, 중간 feature들은 불필요한 정보들이 포함되어있기 때문입니다.)
저자들은 단순한 feature vector의 거리가 아니라, QKVself-attention 구조의 정보를 활용하려고 합니다.
구체적으로, key(K)는 이미 토큰이 담고있는 정보를 dot-product 유사도 계산에 적합한 형태로 만들어주고 있습니다.
그렇기 때문에, 각 토큰의 key들 사이에서 유사도(cosine similarity)를 측정하여, 어떤 토큰들이 비슷한 정보를 담는지 판단합니다.
Token Merging - Bipartite Soft Matching
이제 위에서 Token Similarity를 계산했으니, 실제로 어떤 토큰끼리 매칭할지를 빠르게 결정해야합니다.
K-means clustering같은 알고리즘이 아니라, 매칭 자체의 시간이 무시할 수 있을 만큼 빠른 알고리즘을 선택합니다.
이를 해결하기 위해 채택한 알고리즘이 Bipartite Soft Matching 입니다.
알고리즘은 다음과 같습니다. (위 Figure 1 참고)
- 토큰들을 크기가 비슷한 두 집합 A, B로 나눕니다.
- A의 토큰을 돌면서, B에서 제일 비슷한 토큰 하나를 찾아둡니다.
- 이제 제일 유사도가 높은 r개의 쌍만 남깁니다.
- 연결된 토큰들을 병합합니다. (평균값)
- 두 집합 A, B를 합칩니다.
- 흠, A B를 잘못나누면 항상 최적의 결과가 나올 것 같지는 않네요.
Token Merging - Tracking Token Size
한번 토큰이 병합되면, 해당 토큰은 더 이상 한개의 입력 패치가 아니게 됩니다. (최소 두개 이상의 이미지 패치를 합쳤으니까요) 이는 softmax attention의 결과에 영향을 줄 수 있다고 합니다.
이를 해결하기 위해서 proportional attention을 제안합니다.
$log$ $s$는 각 토큰이 몇개의 패치를 대표하고 있는지를 담은 내용이라고 합니다. (토큰의 개수만큼 정보를 담고 있는 vector s)
많이 더 해질수록 해당 토큰의 attention 계산에 가중치를 더해줍니다.
이 s값은 실제 병합할때도 활용하여 합칩니다.
A토큰과 B토큰을 단순하게 평균낼 때, B는 100개의 토큰을 합친 상태임에도 단순하게 (A+B)/2 로 계산하면 이상하기 때문입니다.
Token Merging - Training with Merging
학습에서 ToMe를 사용하기 위해서, token merging을 일종의 pooling 연산처럼 취급해서 average pooling처럼 역전파를 수행한다고 합니다.
일반적인 ViT 학습 설정으로 ToMe를 그대로 적용해도 최적의 성능을 보였어서, 단순하게 학습 속도를 높이기 위한 옵션으로 사용해도 괜찮을 것 같다고 이야기합니다.
Experiments
