【ICLR2022】Expediting vision transformers via token reorganization
【ICLR2022】Not all patches are what you need: Expediting vision transformers via token reorganization
一些個人想法
- 作者在第5、9、13層應用該方法，這個貌似沒有給出原因
- 圖1用的不是特別好，沒看明白具體要說明什麼
- 該方法和DynamicVIT是最相關的，最近有一系列相關的方法提出，值得關注
- 審稿人指出這個方法難以和swin等多尺度方法結合，因此 the real gain could be very limited
Abstract
之前的VIT將影象塊作為token，來建立多頭自注意力。但是與傳統CNN相比，VIT的 global self-attention between image tokens and long-range dependency 使得模型收斂較慢。目前很少有人關注VIT加速，因為和CNN的差異，CNN加速的一些方法（蒸餾和修建等操作）無法直接應用於VIT。在本文中，作者提出在推理時，重新對token進行組織，實現方法如下：
- 保留關注的影象tokens和融合不關注的影象tokens來重組tokens，以加快後續MHSA和FFN計算，通過這種方式，隨著網路的深入，逐漸減少影象tokens的數量，以降低計算成本
- 在相同數量的輸入影象token情況下，該方法減少了MHSA和FFD計算量，實現高效推理
主要方法
在VIT中，最終是通過 [CLS] 來進行分類，和其他tokens的互動通過encoder的注意力機制完成：
\[\textbf{x}_{class}=\textrm{Softmax}\left(\frac{q_{class}\cdot V}{\sqrt{d}}\right)V=\textbf{a}\cdot V \]這裡的\(q_{class}\)表示[CLS]，組合係數\(\textbf{a}\)是[CLS] 對所有 tokens 的注意力值。因此，可以用該注意力值來表示 tokens 的重要性。
因此作者使用 [CLS] 來識別 tokens 的重要程度，基於這些引數來移除注意力值最小的 token，但從下面的實驗中可以看出，這會嚴重影響分類精度，所以作者在訓練階段加入了影象 token 重組。
作者的總體思路如下圖所示：保留\(k\)個注意力值最大的 tokens，叫做 attentive tokens，然後把其它的 tokens 合併為一個新的 token。
對於 inattentive token，雖然包含較少的資訊，作者認為它們仍然是有助於預測結果的，同時，作者是選擇保留了固定數量的tokens，所以當影象中的物件較大時，移除與之相關的部分，會對效能造成負面影響。綜上，作者將不重要的tokens進行融合（加權平均），融合後將繼續向後傳播 。
實驗分析
實現細節
本文的 token identification module 被新增到 DeiT-S 和 DeiT-B 的第4、7、10層中，以及 LV -ViT-S 的第5、9、13層。另外作者採用了 warm up 的策略來進行 tokens 的標記。讓 attention tokens 的保留率由1逐漸降低到目標值
用高分辨影象訓練
這裡因為對tokens進行了融合減少，所以在維持相同計算量的情況下，允許更多地tokens輸入，因此該方法可以輸入更高解析度影象。作者將標準輸入影象 224×224 調整為 256×256
視覺化
為了進一步研究 EViT 的可解釋性，作者進行了視覺化實驗，可以看到基本上重要的目標區域都得到了保留。