https://arxiv.org/abs/2210.09461
Token Merging: Your ViT But Faster
We introduce Token Merging (ToMe), a simple method to increase the throughput of existing ViT models without needing to train. ToMe gradually combines similar tokens in a transformer using a general and light-weight matching algorithm that is as fast as pr
arxiv.org
Abstract.
본 모델에서는 추가적인 training 없이도 기존 ViT 모델들의 throughput을 증가시키는 간단한 방법인 Token Merging, ToMe을 제안합니다. ToMe는 pruning만큼 빠르면서도 훨씬 더 정확한, 범용적이고 가벼운 매칭 알고리즘을 사용하여 transformer 내에서 유사한 토큰들을 점진적으로 conbines합니다. 지존의 모델에 바로 적용했을 때, ToMe는 이미지 데이터에서 SOTA 모델인 ViT-L @ 512 및 ViT-H @ 518의 처리량을 2배 향상사키고, video data에서 ViT-L의 처리량을 2.2배 향상시키면서도 각 경우의 정확도 하락을 단 0.2% ~ 0.3%로 억제할 수 있습니다.
또한 ToMe는 training 중에도 쉽게 적용할 수 있으며, 실제 비디오 데이터에 대한 MAE(Masked Autoencoders) finetuning 시 훈련 속도를 최대 2배까지 향상시킵니다. ToMe를 적용하여 훈련을 진행하면 정확도 하락을 더욱 최소화할 수 있으며, 오디오 데이터에서 ViT-B 모델의 처리량을 2배로 높이면서도 mAP하락은 단 0.4%에 불과했습니다.
정성적으로 분석한 결과, 본 논문의 저자들은 ToMe가 video의 여러 frames에 걸쳐서도 객체의 여러 부분(parts)들을 하나의 토큰으로 병합해 낸다는 것을 발견했습니다. 전반적으로 ToMe가 보여주는 정확도와 속도는 이미지, 비디오, 그리고 오디오 분야의 SOTA 기술들과 충분히 경쟁할 만한 수준입니다.
*Note: 기존의 token pruning은 배경 이미지(하늘, 잔디 등)처럼 덜 중요해 보이는 토큰들을 점수를 매겨서 아예 drop해 버리는 방식입니다. 이는 빠르긴 하지만, 잘못 판단하면 중요한 정보가 영구적으로 날아가 버려서 정확도가 크게 떨어집니다. 하지만 ToMe의 경우는 버려지는 정보 없이 정보가 압축만 되기 때문에 pruning만큼 연산량은 획기적으로 줄어들면서도 정확도는 거의 그대로 유지됩니다.
Introduction.
Dosovitskiy et al.의 ViTs를 통해 transformer가 NLP에서 Computer Vision 분야로 도입되면서, Computer Vision 분야는 급격히 발전했습니다. 그러나 NLP와 달리, 그 이후 Computer Vision 분야는 vision specific attetnion을 사용하는 Swin, vision specific pooling을 사용하는 MViT, 또는 vision specific conv 모듈을 사용하는 LeViT와 같은 도메인 특화 transformer 하이브리드 모델들에 의해 지배되어 왔습니다. 이러한 추세의 이유는 효율성입니다. vision specific inductive biases들을 추가하면 transformer 모델이 더 적은 연산량으로 더 나은 성능을 낼 수 있습니다.
그럼에도 불구하고 vanilla ViT는 여전히 많은 바람직한 특성을 가지고 있습니다. 단순한 행렬 곱셈으로 구성되어 있어 단순 연산량이 보여주는 것보다 실제 속도가 더 빠릅니다(GPU 가속 덕분에). 또한, MAR와 같이 학습 속도가 빠르면서도 SOTA의 결과를 낼 수 있는 강력한 self-supervised pre-training 기법을 지원합니다. 데이터에 대한 사전 가정이 없기 때문에 다양한 모달리티 전반에 걸쳐 거의 혹은 전혀 수정 없이 적용될 수 있으며, 방대한 양의 데이터에 대해 확장이 용이합니다.
*Note: MAE는 사람이 정답(label)을 알려주지 않아도 인공지능이 스스로 데이터의 특징을 학습하는 방법입니다. 이는 이미지 patch의 75%를 masking한 후 모델에게 남은 25%의 patch만 보여준 뒤, 가려진 75%의 원래 이미지가 무엇이었을 지 복원하게 만듭니다. transformer의 연산량은 입력 토큰 개수에 비례하기 때문에, 전체 이미지의 25%만을 입력으로 넣는 MAE 방식은 학습 속도가 매우 빠릅니다. 한편, inductive bias(사전 가정)이 없다는 것은 무엇일까요? 기존 Computer Vision에서 사용하던 CNN은 이본적으로 서로 가까이 있는 픽셀들끼리는 연관성이 높다는 사전 가정을 가지고 있습니다. 하지만 ViT는 이러한 가정이 없이 이미지를 그냥 1차원으로 길게 늘어놓은 단어들로 취급하기 때문에, 도메인에 맞게 모델을 수정할 필요가 없는 범용성과 일반성을 갖게 됩니다.
그러나 이러한 거대한 모델들을 실햄하는 것은 번거로울 수 있으며, 더 빠른 아키텍처로 큰 모델들의 결과들을 재현하는 것은 어렵습니다. 이에 입력에 구애받지 않는 transformer의 특성 덕분에 runtime에 토큰을 pruning하여 모델을 더 빠르게 만드는 방법들이 각광받고 있습니다. 하지만 token pruning은 여러 단점을 가지고 있는데:
- pruning으로 인한 정보 손실은 합리적으로 줄일 수 있는 토큰 수를 제한합니다.
- 현재의 방법들은 효과를 보기 위해 모델을 재학습해야 합니다(일부는 추가 파라미터 필요).
- 대부분은 학습 속도를 높이는 데 적용할 수 없습니다.
- 몇몇 방법들은 입력 내용에 따라 pruning하는 토큰의 수가 달라지기 때문에 batched inference를 불가능하게 만듭니다.
*Note: token pruning은 추론 속도를 높여주는 반면 학습 속도를 높여주지 못한다고 하는데, 이게 무슨 뜻일까요? GPU는 수십 장의 이미지를 하나의 배치로 병렬 계산할 때 최고 속도를 냅니다. 하지만 pruning은 "A사진은 50개 지우고, B사진은 10만 지워"와 같은 식으로 사진 마다 각각 토큰을 pruning합니다. 즉, 남은 토큰 개수가 달라지게 되므로 묶어서 계산할 수 없는 것입니다. 따라서 이렇게 batch 연산이 깨지는 것을 막기 위해, 기존 모델들은 학습할 때 토큰을 진짜로 지우지 않고 계산은 다 해놓고 결과만 0으로 가려버리는 masking을 사용합니다. 즉, 연산 자체는 100% 돌아가고 있으니 학습 속도는 빨라지지 않는 것입니다(한편 추론은 빛단위로 처리할 필요가 없기 때문에, masking을 하면서까지 batch 단위 계산을 유지할 필요가 없습니다. 따라서 추론 속도는 학습 속도와 달리 빨라지는 것입니다. 추론에서도 배치로 처리해야하는 상황(기업)에서는 똑같은 병목이 존재합니다).
*Note: 한편 ToMe는 토큰을 아예 하나의 토큰으로 압축(merging)해버리고, 이미지마다 줄어드는 토큰의 개수도 정확히 맞추기 때문에, 배치 연산도 잘 돌아가며, 학습 시 역전파를 해야 할 텐서의 크기 자체가 물리적으로 줄어들기 때문에 추론 속도뿐만 아니라 학습 속도까지 줄어드는 것입니다.
*Note: GPU 가속을 위해 이미지마다 압축하는 정도(압축률)를 똑같이 유지한다면, 정보가 많은 이미지는 토큰 압축률을 낮추고, 정보가 적은 이미지는 토큰 압축률을 높여야 한다는 직관에 반합니다. ToMe는 이런 단점에도 불구하고, GPU 가속을 선택했으며, 그나마 drop이 아닌 merging을 통해 정보를 버리지 않음으로써 어느정도 극복했습니다.
본 논문에서는 토큰은 pruning하는 대신 결합하는 Token Merging, ToMe를 제안합니다. 본 논문에서는 맞춤형 매칭 알고리즘 덕분에 이 방법은 pruning 만큼 빠르면서도 더 정확합니다. 더욱이 본 논문의 방법은 학습 유뮤와 관계없이 작동하므로, 정확도 하락을 최소화하면서 거대한 모델에 사용하는 것을 가능하게 합니다. 학습 중에 ToMe를 사용하면 실제 학습 속도가 증가하는 것을 관찰할 수 있으며, 일부 경우에는 전체 학습 시간을 절반으로 단축시킵니다. 또한 이미지, 비디오, 오디오에 아무런 수정 없이 ToMe를 적용하였고, 모든 경우에서 SOTA와 경쟁력이 있습니다.
Related Work.
Efficient Transformers.
NLP와 vision 분야 모두에서 더 효율적인 transformer를 만들기 위한 여러 연구가 시도되었으며, 일부는 더 빠른 attention에 중점을 두고(Choromanski et al., 2020; Dao et al., 2022), 일부는 heads나 features들을 pruning하려 시도하며(Wang et al., 2022; Voita et al., 2019), 일부는 도메인 특화 모듈을 주입하려고 시도합니다(Mehta & Rastegari, 2021). 본 논문에서는 때로는 추가 학습 없이도 토큰을 병합하여 기존 ViR 모델의 속도를 높임으로써 더 복잡한 도메인 특화 모델들의 speed-accuracy trade off 의 최적을 달성하는 데 초점을 맞춥니다.
Token Reduction.
transformer는 임의의 개수의 토큰으로 작동할 수 있기 때문에, 최근 여러 연구가 NLP와 vision 분야 모두에서 transformer의 token을 pruning하려 시도했습니다. 그러나 이러한 방법들은 학습이 필요한 반면, 본 논문에서 제안한 방법은 학습 없이도 사용할 수 있습니다. 더욱이 대부분의 pruning 연구들은 dynamic입니다. 즉, 이미지나 문장마다 토큰마다 토큰의 수가 달라집니다. 이는 정확도에는 이점이 되지만, 토큰 수가 다른 샘플들을 더 이상 batch 단위로 처리할 수 없게 되므로 실용성을 제한합니다. 이 문제를 해결하기 위해 대부분의 pruning 논문들은 학습 중에 토큰을 제거하는 대신 mask를 적용하는데, 이는 pruning을 통한 속도 향상 효과를 무효화시킵니다. 반면 본 논문의 방식은 추론과 학습 모두에 적용될 수 있으며, 두 경우 모두에서 실제 속도 향상을 달성합니다.
Token Merging.
Strategy.
Transformer의 각 block에서, 토큰을 merging하여 layer당 r개씩 토큰을 줄입니다. r은 토큰의 비율(ratio)가 아니라 수량(quantity)입니다. 네트워크 내의 L개 블록에 걸쳐, 점진적으로 rL개의 토큰을 병합합니다. r값을 변화시키면 속도와 정확도 간의 trade-off가 발생합니다. 토큰이 적어지면 정확도는 낮아지지만 처리량은 높아지기 때문입니다. 중요한 점은, 이미지의 내용과 무관하게 rL개의 토큰을 고정적으로 줄인다는 것입니다. 일부 pruning 방법들은 token의 수를 동적으로 변화시킵니다. 이는 정확도를 높여주지만, padding token 없이는 batched inference이나 학습을 불가능하게 하므로 일반적으로 실용적이지 않습니다.
위 그림 1. 에서 볼 수 있듯이, 각 transformer block의 attention과 MLP branch 사이에 token merging 단계를 적용합니다. 이는 토큰 축소 기법을 블록의 시작 부분에 배치하는 경향이 있는 기존 연구들과 대조되는 점입니다. 해당 배치는 병합될 토큰들로부터 정보가 전파될 수 있게 하며, attention 내의 features을 사용하여 무엇을 병합할지 결정할 수 있게 해 주는데, 이 두 가지 모두 정확도를 향상시킵니다.
*Note: 기존의 논문들은 보통 transformer block의 맨 처음(attention을 하기 전)에 토큰을 줄였습니다. 하지만 ToMe는 attention과 MLP 사이로 merging module을 추가했는데, 이는 다음과 같은 이점이 존재합니다:
- Imformation Propagation 보존
- Attention은 토큰들끼리의 관계에 대해 파악하고 서로 정보를 교환하는 작업입니다. 만약, attention 이전에 토큰 A와 B를 A+B로 뭉쳐버리면, 다른 토큰(C, D, E)들은 A와 B의 섬세하고 개별적인 특징을 보지 못하고 뭉뜽그려진 정보만 보게 됩니다.
- 하지만 attention 이후 merging을 수행하면, 토큰 C, D, E는 attention 과정에서 이미 A와 B의 원본 정보를 충분히 참고하고 자신의 값을 업데이트했습니다. 즉, merging 되어 사라질 토큰들의 정보가 이미 네트워크 전반에 안전하게 전파된 후 merging이 일어나므로 정보 손실이 줄어드는 것입니다.
- Attention feature(keys)의 재활용
- 토큰을 merging하려면 누구랑 누가 비슷한가?를 계산해야 합니다. ToMe는 유사도 계산을 위해 추가 연산을 하지 않고, attention 과정에서 이미 만들어진 K vector를 그대로 가져다 씁니다. key vector 자체가 '이 토큰이 어떤 정보를 담고 있는지'를 잘 용약한 지표이기 때문입니다.
- Attention을 하기 전에 merging을 시도하면 아직 이번 layer의 key vector가 만들어지지 않았으므로 이 정보를 사용할 수 없습니다. 즉, attention 직후에 merging module을 두면 정확도 높은 유사도 판단 기준을 연산량 추가 없이 덤으로 얻게 되는 것입니다.
Token Similarity.
유사한 토큰들을 병합하기 전에, 먼저 "유사하다"는 것이 무엇을 의미하는지 정의해야 합니다. 두 토큰 간의 feature 거리가 짧으면 유사하다고 부르고 싶을 수 있지만, 이것이 반드시 최적인 것은 아닙니다(Marin et al., 2021). 최신 transformer의 intermediate feature space는 과대적합되어 있습니다. 예를 들어, ViT-B/16은 각 토큰의 RGB 픽셀 값 전체를 인코딩하기에 충분한 feature를 가지고 있습니다. 이는 intermediate features들이 유사도 계산을 교란할 수 있는 무의미한 노이즈를 포함할 가능성이 있음을 의미합니다.
다행히도 transformer는 QKV self-attention을 통해 이 문제를 자체적으로 해결합니다. 구체적으로, K는 dot product similarity 계산에 사용될 수 있도록 각 token에 포함된 정보를 이미 요약하고 있습니다. 따라서 각 토큰의 ket 간의 내적 유사도 지표(예: cosine similarity)를 사용하여 어떤 토큰들이 유사한 정보를 포함하고 있는지 결정합니다(아래 표 1.(a), 1.(b) 참조).
*Note: Attention 연산 자체가 feature들의 dot product를 기준으로 뭉치는 작업을 수행하며, 이 정보를 Key가 가지고 있기 때문에, key들에 대한 cos sim을 통해서 유사도를 정의할 수 있다는 뜻이다.
Bipartite Soft Matching.
토큰 유사도가 정의되었으므로, 전체 수를 r만큼 줄이기 위해 어떤 토큰을 매칭할지 결정하는 빠른 방법이 필요합니다. 이 문제에 대한 몇 가지 잠재적인 해결책으로 K-means clustering이나 graph cuts 등이 있습니다. 하지만 해당 모델에서는 수천 개에 달할 수 있는 토큰들에 대해 이 매칭을 L번이나 수행해야 하므로, 그 실행 시간은 절대적으로 무시할 수 있을 만큼 짧아야 합니다. 대부분의 반복적인 군집화 알고리즘은 여기에 부합하지 않습니다.
따라서 더 효율적인 해결책을 본 논문에서는 제안합니다. 목표는 다음과 같습니다: 1) 병렬화할 수 없는 반복적인 모든 작업을 피한다. 2) merging으로 인한 변화가 gradual이어야 한다. 2)이 군집화(clustering)이 아닌 matching에 초점을 맞추는 이유입니다. 군집화는 하나의 그룹으로 병합될 수 있는 토큰의 수에 제한을 두지 않는 반면, 매칭은 대부분의 토큰을 병합되지 않은 상태로 유지합니다. 본 논문에서 제안한 알고리즘은 다음과 같습니다(위 그림 1. 에 시각화 됨):
- 토큰들을 크기가 거의 같은 두 집합 A와 B로 분할합니다.
- A의 각 토큰에서 B에 있는 가장 유사한 토큰으로 하나의 edge를 연결합니다.
- 가장 유사한 r개의 edges만 유지합니다.
- 여전히 연결되어 있는 토큰들을 병합합니다(예: mean feature).
- 두 집합을 다시 하나로 이어 붙입니다(concatenate).
*Note: merging으로 인한 변화가 점진적(gradual)이어야 한다는 것은 한 layer를 통과할 때 전체 토큰 수가 얼마나 크게 요동치지 말아야 한다는 것입니다. 군집화를 사용하면 모든 토큰이 무조건 어떤 그룹이든 소속되어야 합니다. 196개의 토큰을 180개로 줄이려고 K-means를 돌리면, 196개 토큰 전체가 180개의 그룹(gradual하게 # tokens를 조절하기 위해)으로 강제 재배치됩니다. 즉, 단 하나의 레이어만 지나도 (강제로 임의의 토큰들이 그룹에 묶이면서) 전체 토큰의 성질이 완전히 뒤섞여버리는 급진적인 변화가 일어납니다. 한편, ToMe가 고안한 bipartite matching 알고리즘은 한 번의 layer에서는 아주 소수의 확실하게 비슷한 녀석들(r개)만 뭉치고, 대다수의 토큰은 원본 그대로 내버려 두겠다는 뜻입니다. 이 과정을 12개 layers에 걸쳐 천천히 반복하면서 최종적으로 r x L 개의 토큰을 줄여나갑니다.
*Note: 이는 추가 학습 없이 기존 모델에 적용해도 성능이 떨어지지 않는 성질과 관련됩니다. 만약 clustering을 통해 토큰들을 한 번에 섞어버리면, 모델의 다음 leyer(MLP)는 평소에 보던 feature value가 아니기 때문에 모델은 당황하게 되며, 모델의 밸런스가 붕괴되어 반드시 재학습(fine-tuning)을 통해서 이 낯선 값들에 모델을 적응시켜야만 합니다. 하지만 ToMe 처럼 gradual으로 줄이면 어떨까요? 196개의 토큰들 중 160개는 평소에 보던 익숙한 원본 토큰 그대로 들어오고, 16개 정도만 살짝 섞인 값이 들어옵니다. 이는 모델 입장에서 아주 미세한 변화이며 치명적인 오차로 받아들여지지 않습니다.
이 과정이 bipartite graph를 생성하고 A의 각 토큰이 단 하나의 edge만 가지기 때문에, 4단계에서 connected components를 찾는 것은 매우 간단합니다. 게다가 모든 토큰 쌍 사이의 유사도를 일일이 계산할 필요가 없으며, A와 B를 신중하게 선택한다면 이는 정확도에 문제가 되지 않습니다. 실제로 이 bipartite soft matching은 토큰을 무작위로 dropping하는 것만큼이나 빠르며, 코드로 구현하는 데 단 몇 줄 밖에 걸리지 않습니다.
*Note: 왜 모든 토큰 쌍 사이의 유사도를 일일이 계산할 필요가 없을까요? 이는 전체 토큰에 대해서 NxN의 유사도를 계산할 필요 없이, 각 집합(N/2)에 대해서 유사도를 구하면 되기 때문에 N^2/4의 유사도만 구하면 된다는 뜻입니다(A는 A끼리 유사도 계산 x). 한편, A와 B를 신중하게 선택한다면 정확도 문제가 없다는데, 어떻게 A와 B를 구할까요? 이는 아래 표 1.(d)에 나와있는 alternating을 사용합니다. 이미지의 경우 가까이 붙어있는 patch(pixel)일수록 내용이 비슷할 확률이 높습니다. 따라서 토큰들을 A와 B로 나눌 때, 사진의 위/아래로 쪼개거나 random으로 쪼개지 않고, 체스판 무늬처럼 교대로(alternating) 배정합니다. 이렇게 배정하면, 가장 비슷할 확률이 높은 바로 옆자리 이웃 patch는 무조건 나와 반대 집합에 배정되기 때문에, merging될 확률이 높습니다.
Tracking Token Size.
토큰들이 병합되고 나면, 더 이상 단일 입력 patch 하나만을 나타내지 않게 됩니다. 이는 softmax attention의 결과를 바꿀 수 있습니다. 만약 동일한 key를 가진 두 토큰을 병합한다면, 그 key는 softmax 항에서 더 적은 영향을 미치게 됩니다. 이를 proportional attention이라는 간단한 수식을 통해 수정할 수 있습니다:
여기서 s는 각 token의 크기(token이 나타내는 원래 patch의 수)를 포함하는 행 벡터(row vector)입니다. 이는 key가 s개 복사되어 있는 것과 완벽히 동일한 연산을 수행합니다. 또한 두 토큰을 하나로 병합할 때처럼 토큰들이 aggregated될 때마다, 토큰의 크기 s를 가중치로 부여해야 합니다.
*Note: 만약 모델이 "지금 파란 하늘 정보가 얼마나 중요해?"라고 물었을 때, 사진 속에 '파란 하늘' 패치가 10개 있다면 10표의 영향력을 행사하게 됩니다. ToMe 적용 후, ToMe가 10개의 토큰이 중복됨을 확인하고 1개의 토큰으로 merging한다면, 다음 layer에서 softmax attention은 이 토큰을 보고 1표의 가중치만을 할당합니다. 원본 사진에서는 10표의 비중을 차지했던 patch가 merging후에는 1/10으로 가중치가 줄어드는 것입니다. 이를 보완하기 위해, merged patch가 원래는 s개의 patch가 합쳐진 것임을 이용해 attention score의 s배를 취하기 위해, log(s)를 더해주는 것입니다.
Training with Merging.
지금까지의 각 구성 요소는 이미 학습된 ViT 모델에 token merging 모듈을 추가할 수 있도록 설계되었습니다. ToMe와 함께 학습을 진행하는 것이 필수적인 것은 아니지만, 정확도 하락을 줄이거나 학습 속도를 높이기 위해 바람직할 수 있습니다. 학습을 진행할 때, 단순히 token merging을 pooling 연산으로 취급하고, 마치 average pooling을 사용하는 것처럼 merged tokens들을 통해 backprop을 수행합니다. token pruning에서 사용하는 gumble softmax와 같은 복잡한 gradient tricks을 사용할 필요성이 없으며, 실제로 vanilla ViT를 학습할 때 사용된 것과 동일한 설정이 여기서도 최적임을 확인했습니다.
Image Experiments.
Design Choices.
위 표 1. 에서 ablation을 수행합니다. 각 ablation마다, 보라색으로 표시된 기본 파라미터에서 시작합니다. 특별히 명시되지 않는 한, 추가 학습 없이 ViT-L/16 MAE 모델에서 테스트하며, 네트워크 24개 layers에 걸쳐 점진적으로 토큰 98%를 제거하는 r = 8로 병합을 수행합니다.
토큰의 feature X는 성능 측면에서 최선이 아닙니다(표 1.(a) 참조). merging 연산을 attention 이후로 옮기고 attention key를 사용하는 것이 훨씬 더 정확합니다. 그 다음, 표 1.(b)에 나타난 바와 같이 토큰 거리를 측정하는 데에는 cosine similarity가 가장 좋습니다. 마지막으로 효율성을 위해 attention heads들에 걸쳐 K를 이어 붙이는(concatenating) 대신 평균을 냅니다(표 1.(c)).
어떤 토큰을 병합할지 결정한 후, 토큰 크기 s로 가중치를 둔 평균을 통해 그것들을 결합합니다. 표 1.(d)에서 이는 B에 있는 토큰만 유지하거나, max pooling 또는 가중치 없는 mean pooling을 사용하는 것 보다 우수한 성능을 보여줍니다. 그 다음, bipartite matching 알고리즘은 입력을 두 개의 서로 겹치지 않는 집합으로 분할할 것을 요구합니다. 분할 후 다시 집합들을 이어 붙이기 때문에, 교대로(alternating) 토큰을 할당하는 것이 가장 잘 작동함을 확인했습니다(표 1.(e)). A를 다 채운 다음 B를 채우는 방식은 가장 낮은 성능을 보였습니다.
merging되고 나면 token은 둘 이상의 입력 patch를 나타낼 수 있습니다. 이는 위에서 살펴봤던 proportional attention을 통해 해결하며, 표 1.(f)에서 이에 대한 ablation을 수행합니다. proportional attention은 supervised trained model(예: AugReg, SWAT, etc.)에는 필수적이지만, MAE 모델에는 그렇지 않습니다. 학습 후에는 이러한 차이가 사라지는 것으로 보아, 이는 MAE가 사전 학습 중에 이미 토큰을 제거하기 때문일 가능성이 큽니다.
위 표 2.에서 bipartite matching을 pruning과 merging을 모두 포함하는 다른 token reduction 알고리즘들과 비교합니다. pruning은 빠르지만, 전반적으로 토큰의 98%가 제거되면 중요한 정보가 소실됩니다. 이는 무작위로 pruning하는 경우와 attention을 받지 못한 것을 기반으로 pruning을 하는 경우(Kim et al., 2021) 모두에 해당합니다. 반면, 토큰을 merging하는 것은 유사하지 않은 토큰들이 merging될 때만 정보를 잃습니다. 따라서 merging할 유사한 토큰들을 올바르게 선택하는 것이 중요합니다.
*Note: attention을 받지 않은 원시 feature를 기반으로 pruning을 수행하면, attention 연산의 결과로 feature들이 얻을 수 있는 tokens들 간의 context를 반영하지 않은 채, drop될 수 있다는 문제가 있습니다. 이는 위 표 1.(a)의 Xpre를 이용해서 merging했을 때 정확도가 낮은 것과 일맥상통합니다.
처음에는 K-means가 분명한 선택처럼 보일 수 있지만, 속도가 느릴 뿐만 아니라 pruning보다 아주 약간만 더 낫습니다. K-means는 reconstruction error를 최소화할 수는 있지만(min||A-Ac||), 다수의 토큰들이 동일한 cluster에 매칭되는 것을 허용하여 유사하지 않은 토큰이 병합될 확률을 높입니다. Marin et al., (2021)은 K-means를 기반으로 한 몇 가지 더 빠른 clustering 알고리즘을 연구하지만, 그들의 설정에서는 학습 없이는 10% 미만의 정확도 하락을 얻어내는 데 실패합니다.
대신, 가장 유사한 토큰들만 병합하는 매칭 알고리즘을 사용하며, 가장 유사한 토큰 쌍을 병합한 다음 복원 추출 없이 이를 r번 반복하는 greedy 방식으로 수행할 수 있습니다. 이는 정확하지만 sequential이어서 r이 클 경우 속도가 느려질 수 있습니다. 따라서 본 논문에서 제안하는 bipartite matching은 이 greedy 알고리즘의 정확도와 pruning의 속도를 모두 가지면서도 r에 대해 일정한 실행 시간을 가집니다.
(... 나머지는 원본 논문 참조 ...)
Visualizations.
위 그림 4.는 네트워크 끝부분에 있는 각 merged token에 속하는 입력 patch들을 보여줍니다.
