【導(dǎo)讀】由Meta AI的研究人員推出Token Merging（ToMe），無需訓(xùn)練即可加速 ViT 模型。更重要的是，這個(gè)模型不需對(duì)token進(jìn)行剪枝。

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

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

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

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

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

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

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

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

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

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

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

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

TokenMerging：換個(gè)想法

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

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

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

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

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

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

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

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

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

具體做法

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

因?yàn)閗ey已經(jīng)總結(jié)了每個(gè)token中包含的信息，以便用于Attention中的dot-product來衡量token間的相似度。

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

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

Meta團(tuán)隊(duì)沒有使用kmeans聚類算法或圖分割算法，而是使用匹配算法，因?yàn)楹笳卟粌H可以精準(zhǔn)匹配每一層token的數(shù)量，還能快速執(zhí)行上千次匹配。這些都是迭代聚類算法無法完成的。

因此，Meta團(tuán)隊(duì)提出了一個(gè)更有效的解決方案。

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

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

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

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

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

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

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

審核編輯 ：李倩