加速ViT模型新思路！ Meta推出Token Merging，不靠剪枝靠合併

視覺變換器（ViT）在兩年前進入大眾視野，並成為電腦視覺研究的核心組成部分。

它成功將一個在自然語言處理領域的Transformer模型遷移到電腦視覺領域。從那時起，電腦視覺領域的進步已經加速。

儘管在成本與性能方面被超越，Vanilla ViT仍有許多優點。

它們是由簡單的矩陣乘法組成的，這使得它們的速度比它們的原始運算量所顯示的要快。

此外，它們支援強大的自監督預訓練技術，如MAE（掩碼自動編碼器），可以產生最先進的結果，同時可以進行快速訓練。

而且由於它們不會對資料進行假設，它們可以幾乎不加改變地應用在圖片、音訊、文字等諸多模式中。

當然，理想很豐滿，現實很骨感。 ViT模型的規模大，有較大延遲。在資源有限的設備上，運行這個複雜模型會產生很大問題。

Token剪枝：變好了，但沒完全好

針對運算慢的問題，研究人員給出了多個解決方案。其中一個常見的加速視覺 Transformer模型的方法是對進行token剪枝。

在運行時修剪標記，透過修剪不太重要的token產生高效的Transformer。如DynamicViT分層修剪冗餘token，從而在分類任務中實現FLOPs減少。

然而，token剪枝有幾個問題，其中最主要的，是由於修​​剪token會產生資訊損失，因此，人們對ViT模型token的剪枝數量是有限的，為了減少資訊損失，只能修剪不重要的token。

而且，為了讓修剪過的token有效，人們需要再次訓練模型。這就造成額外的資源消耗。

更重要的是，token剪枝是動態的過程，需要根據不同的圖像或句子確定token剪枝的不同數量。雖然這有利於提高準確性，但卻不夠實用實用性，因為這種情況下，資料就不能再進行批次處理。

為了解決這個問題，人們需要在剪枝過程中加入掩碼，而這會進一步影響效率的提升。

簡單來說，token剪枝確實讓ViT跑得更快，但這是在資訊損耗的代價上實現的。

#TokenMerging：換個想法

怎麼才能使ViT的速度類似剪枝，但保持比剪枝更高的準確度呢？ Meta AI研究團隊給了新的解題想法：Token Merging（ToMe）。

論文連結：https://arxiv.org/pdf/2210.09461.pdf

#Token Merging選擇將token結合，而非進行剪枝。由於其定制的匹配演算法，它和剪枝一樣快，同時更準確。另外，它的工作不需要任何額外的訓練，所以你可以在巨大的模型上使用它來加快它們的速度，而不會犧牲很多準確性。

Meta的目標是在現有的ViT中插入一個Token Merging的模組，透過合併冗餘的token，在不需要額外訓練的前提下提高訓練和推理的吞吐量。

基本想法是：在Transformer模型中，透過合併，使每層減少r個token。假設一個Transformer模型有L層，那麼透過合併就可以減少rL個token。變數r的大小決定了速度和精確度的關係，因為更少的標記意味著更低的準確度但更高的吞吐量。

值得注意的是，在Token Merging中，無論圖片的內容為何，都會減少rL標記。這完美解決了token剪枝中無法進行批次處理的問題。

透過ToMe，類似的token批在每個Transformer區塊中合併：例如，狗的皮毛被合併成一個token。

Token Merging被插入每個attention區塊和每個Transformer區塊。這也與token剪枝的工作流程形成對比。後者傾向於將剪枝步驟放在每個Transformer區塊的開頭。

透過Token Merging，需要合併的token的資訊可以傳播，ViT也能夠藉助attention區塊中的特徵來決定需要合併哪些token。

具體做法

##合併的第一步是確定相似的token。在Transformer中的QKV（query, key, value）已被提取的條件下，透過消融實驗，研究團隊發現使用key可以最好衡量token之間的相似度（下圖紫色部分）。

因為key已經總結了每個token中包含的信息，以便用於Attention中的dot-product來衡量token間的相似度。

除了研究哪個指標比較好衡量token相似度外，還需要知道什麼距離衡量相似度。透過實驗研究團隊發現，使用使用餘弦距離來衡量toke之間的相似度可以獲得最佳的精確度和速度的關係。

確定了token的相似性，接下來需要一個快速的方法來確定哪些token需要匹配，以減少總數的r。

Meta團隊沒有使用kmeans聚類演算法或圖分割演算法，而是使用匹配演算法，因為後者不僅可以精確地匹配每一層token的數量，還能快速執行上千次匹配。這些都是迭代聚類演算法無法完成的。

因此，Meta團隊提出了一個更有效的解決方案。

設計目標如下。 1.)避免任何無法並行化的迭代，2.)希望合併的變化是漸進的，因為聚類對多少個標記可以合併到一個組中沒有限制（這可能會對網絡產生不利影響），而匹配則使大多數標記沒有被合併。 #

1. 將所有token分成相同大小的2個集合A與B。
2. 將從集合A中的每個token到B中與其最相似的token畫一條邊。
3. 只留下最相似的r條邊, 其餘刪掉。
4. 融合仍然相連的邊（特徵取平均值）。
5. 把這兩個集合拼在一起, 得到最終的合併結果。

透過這項獨特的技術，可以提升ViT模型的吞吐量和實際訓練速度。使用Token Merging可以將訓練速度提高一倍。它可以用於圖像、視訊和音訊任務，並且仍然可以達到最先進的準確性。

#

以上是加速ViT模型新思路！ Meta推出Token Merging，不靠剪枝靠合併的詳細內容。更多資訊請關注PHP中文網其他相關文章！