# 【論文筆記】Token Merging: Your ViT But Faster

論文連結：
https://arxiv.org/abs/2210.09461
發表於 ICLR 2023

## Introduction
這篇研究介紹了 Token Merging (ToMe)，是一個用於提升 ViT 模型 throughput 的方法。這個方法可以用於 training，提升訓練的速度，也可以直接使用在訓練好的模型上，提升 inference time 的 throughput。

ToMe 的基本概念是透過 matching algorithm 將相似的 tokens 合併在一起，讓 transformer 可以不用對這麼多 tokens 進行運算，因而提升速度。實驗發現 ToMe 的準確度和速度都能和 state-of-the-art 相匹敵。

## Token Merging
### Strategy
採用的策略是在 transformer 的每一個 block，都進行 token merging 來減少 $r$ 個 tokens，因此經過 $L$ 個 blocks 之後，就會一步步減少總共 $rL$ 個 tokens。$r$ 是一個可以調整個參數，不同的 $r$ 會產生 speed-accuracy 之間 trade-off，比較高的 $r$ 代表比較少的 tokens，因此會讓 accuracy 降低，但 throughput 比較高。

如下圖 (b) 所示，每個 transformer block 中，token merging 是在 attention 和 MLP 之間進行。作者指出這樣的設計有兩個特點，一是就算是可能合併的 tokens，訊息也可能在它們之間傳播，二是我們可以使用 attention 當中的 feature 來決定我們要合併哪些 tokens。

### Token Similarity
要合併相似的 tokens，就需要找出哪些 tokens 是相似的。作者指出說大家可能會直觀認為計算 features 之間的 distance 就可以了，但事實上 transformer 中的 intermediate feature space 是非常 overparameterized（過度參數化）的，代表說 intermediate features 很有可能含有很多不重要的資訊，讓相似度的計算不準確。這裡他們提出使用 QKV self-attention 當中的 key (K) 作為代表，計算兩個 keys 之間的 cosine similarity，來決定兩個 tokens 是否具有相似的資訊。

### Bipartite Soft Matching
為了要進行加速，這裡需要一個很快速的方法來決定哪些 tokens 是相似的。作者提出了 bipartite soft matching，如上圖 \(c\) 所示，步驟如下：
1. 將所有的 tokens 大約平分成 A 和 B 兩個 sets
2. 每一個 set A 的 token 都用一條邊連接到一個在 set B 最相似的 token
3. 保留 $r$ 個最相似的邊
4. 合併 tokens，例如計算兩個 features 的平均
5. 把兩個 sets 再整合回一個 set

整個步驟不需要任何迭代，可以平行進行運算，因此非常快速。另外雖然並沒有計算每一對 tokens 之間的相似度來精準找出最像的 tokens，但只要選擇好的方式分出 set A 和 B，就不會對最後的準確度造成問題。

### Tracking Token Size
最後還會遇到的另一個問題是當我們將 tokens 合併之後，一個 token 就不會只代表一個 patch，這會影響到之後 attention 的計算，因此作者提出了 proportional attention 來解決這個問題，重新定義 attention 矩陣 $A$ 的計算為
$$
A = \text{softmax} \Bigr(\frac{QK^T}{\sqrt{d}} + \log s \Bigr)
$$

其中 $s$ 會是一個 row vector，表示各個 tokens 所代表的 patch 數量。這個操作可以調整 attention 的權重分配。

## Experiments
### Design Choice
作者在 ImageNet-1k dataset 上進行實驗。下圖顯示他們在各種設計上進行 ablation study 的結果，測試的模型是 ViT-L/16 MAE，沒有再次 training，而是直接進行 inference，$r$ 則是設為 8，紫色註記代表預設使用的設定：


(a) 選擇 key 作為計算相似度的 feature 是因為可以產生最高的準確度
(b) 可以看到用 cosine similarity 計算相似度所產生的結果可以有最高的相似度
(c\) 合併 features 時選擇用平均的方式雖然準確度比 concat 低，但速度比較快
(d) 計算 features 的平均有很多方式，實驗發現用 size $s$ 計算加權平均的準確度最高
(e) 因為之後重新整合 sets 是用 concatenate 的方式，分配 sets 的時候使用交替分配的方式可以得到最高的準確度
(f) Proportional attention 在 MAE 模型上是不需要的，但在其他 ViT 模型上（例如 AugReg）則有改善準確度的效果，猜測可能的原因是 MAE 模型在訓練的時候就已經 mask 部分的 tokens


上表顯示在相同條件下，不同 matching 演算法所得到的結果，可以看到雖然 bipartite matching 所得到的準確度會比 greedy matching 低一點點，但 throughput 上有比較好的表現。

### Model Sweep

為了了解 $r$ 的影響，下面的實驗將 token merging 應用在各種 SOTA ViT 模型上，針對每一個模型，從 $r=0$ 測試到 $r$ 太大無法合併 tokens 為止，構建出 throughput-accuracy 的 trade-off 曲線。

AugReg 和 SWAG 是兩種 supervised models，從下圖可以看到隨著 $r$ 增加，可以有接近 2 倍的 thoughput，越大的模型準確度則下降越少，最大的模型甚至幾乎沒有下降。作者推測是因為大的模型深度比較深，可以讓 merging 相對比較平緩，降低 merging 造成的影響。

MAE 是一個 self-supervised 模型，因此如果直接套用 token merging 而不重新訓練，準確度降低的幅度會比較大，但如果使用 token merging 並重新 fine-tune 模型的話，準確度降低的幅度會比較小。


### Comparisons

下表比較了加上 token merging 的 MAE fine-tuned model 和其他 SOTA models，都只訓練在 ImageNet-1k。作者指出加上 ToMe，就可以允許使用更大的模型產生同樣量級的 throughput，例如 ViT-H 加上 ToMe 之後，throughput 就可以和原始的 ViT-L 在相同的量級。


### Visualization

下圖顯示在模型的最後，每一個 patch 是如何 merged 在一起的。可以看到利用 ToMe，不論在 foreground 或 background 的 tokens 都可以被合併，不會損失資訊。


## Conclusion

這篇研究介紹了 Token Merging (ToMe)，利用逐步合併 tokens 的方法來增加 ViT 模型的 throughput。雖然這篇研究主要專注在分類問題上，但作者認為 ToMe 是可以運用到各式各樣的任務上的。