加速ViT模型新思路！Meta推出Token Merging

【導讀】由Meta AI的研究人員推出Token Merging（ToMe），無需訓練即可加速 ViT 模型。更重要的是，這個模型不需對token進行剪枝。

視覺變換器（ViT）在兩年前進入大眾視野，并成為計算機視覺研究的核心組成部分。 它成功將一個在自然語言處理領(lǐng)域的Transformer模型遷移到計算機視覺領(lǐng)域。從那時起，計算機視覺領(lǐng)域的進步已經(jīng)加速。

盡管在成本與性能方面被超越，Vanilla ViT仍有許多優(yōu)點。

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

此外，它們支持強大的自監(jiān)督預訓練技術(shù)，如MAE（掩碼自動編碼器），可以產(chǎn)生最先進的結(jié)果，同時可以進行快速訓練。

而且由于它們不對數(shù)據(jù)進行假設(shè)，它們可以幾乎不加改變地應(yīng)用在圖片、音頻、文本等諸多模式中。

當然，理想很豐滿，現(xiàn)實很骨感。ViT模型的規(guī)模大，有較大延時。在資源有限的設(shè)備上，運行這個復雜模型會產(chǎn)生很大問題。

Token剪枝：變好了，但沒完全好 針對運算慢的問題，研究人員給出了多個解決方案。其中一種常見的加速視覺 Transformer模型的方法是對進行token剪枝。 在運行時修剪標記，通過修剪不太重要的token產(chǎn)生高效的Transformer。如DynamicViT分層修剪冗余token，從而在分類任務(wù)中實現(xiàn)FLOPs減少。

然而，token剪枝有幾個問題，其中最主要的，是由于修剪token會產(chǎn)生信息損失，因此，人們對ViT模型token的剪枝數(shù)量是有限的，為了減少信息損失，只能對不重要的token進行修剪。

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

更重要的是，token剪枝是動態(tài)的過程，需要根據(jù)不同的圖像或句子確定token剪枝的不同數(shù)量。雖然這有利于提高準確性，但卻不夠?qū)嵱脤嵱眯裕驗檫@種情況下，數(shù)據(jù)不能再進行批處理。

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

簡單來說，token剪枝確實讓ViT跑得更快，但這是在信息損耗的代價上實現(xiàn)的。

TokenMerging：換個想法

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

論文鏈接：https://arxiv.org/pdf/2210.09461.pdf

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

Meta的目標是在現(xiàn)有的ViT中插入一個Token Merging的模塊，通過合并冗余的token，在不需要額外訓練的前提下提高訓練和推理的吞吐量。

基本思路是：在Transformer模型中，通過合并，使每層減少r個token。假設(shè)一個Transformer模型有L層，那么通過合并就可以減少rL個token。變量r的大小決定了速度和精度的關(guān)系，因為更少的標記意味著更低的準確度但更高的吞吐量。

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

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

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

通過Token Merging，需要被合并的token的信息可以得到傳播，ViT也能夠借助attention塊中的特征來決定需要合并哪些token。

具體做法

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

因為key已經(jīng)總結(jié)了每個token中包含的信息，以便用于Attention中的dot-product來衡量token間的相似度。

除了研究哪個指標更好衡量token相似度外，還需要知道什么距離衡量相似度。通過實驗研究團隊發(fā)現(xiàn)，使用使用余弦距離來衡量toke之間的相似度可以獲得最好的精度和速度的關(guān)系。

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

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

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

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

將所有token分為相同大小的2個集合A與B。

把從集合A中的每個token到B中與其最相似的token畫一條邊。

只留下最相似的r條邊, 其余刪掉。

融合仍然相連的邊（特征取均值）。

把這兩個集合拼在一起, 得到最終的合并結(jié)果。

通過這項獨特的技術(shù)，可以提高ViT模型的吞吐量和實際訓練速度。使用Token Merging可以將訓練速度提高一倍。它可以用于圖像、視頻和音頻任務(wù)，并且仍然可以達到最先進的準確性。

審核編輯 ：李倩