DVT：華為提出動態級聯Vision Transformer，效能槓槓的 | NeurIPS 2021

語言: CN / TW / HK

時間 2022-09-19 09:32:52 曉飛的演算法工程筆記

主題: Transformer NeurIPS

論文主要處理Vision Transformer中的效能問題，採用推理速度不同的級聯模型進行速度優化，搭配層級間的特徵複用和自注意力關係複用來提升準確率。從實驗結果來看，效能提升不錯

來源：曉飛的演算法工程筆記公眾號

論文: Not All Images are Worth 16x16 Words: Dynamic Transformers for Efficient Image Recognition

論文地址：http://arxiv.org/abs/2105.15075
論文程式碼：http://github.com/blackfeather-wang/Dynamic-Vision-Transformer

Introduction

Transformers是自然語言處理 (NLP) 中占主導地位的自注意的模型，最近很多研究將其成功適配到影象識別任務。這類模型不僅在ImageNet上取得了SOTA，而且效能還能隨著資料集規模的增長而不斷增長。這類模型一般都先將影象拆分為固定數量的影象塊，然後轉換為1D token作為輸入，拆分更多的token有助於提高預測的準確性，但也會帶來鉅額的計算成本（與token數成二次增長）。為了權衡效能和準確率，現有的這類模型都採用14x14或16x16的token數量。

論文認為不同圖片之間存在相當大的差異，使用相同數量的token處理所有圖片並不是最優的。最理想的做法應為每個輸入專門配置token數量，這也是模型計算效率的關鍵。以T2T-ViT-12為例，官方推薦的14x14 token數僅比4x4 token數增加了15.9%（76.7% 對 60.8%）的準確率，卻增加了8.5倍的計算成本（1.78G 對 0.21G）。也就是說，對“簡單”圖片使用14x14 token數配置浪費了大量計算資源，使用4x4 token數配置就足夠了。

受此啟發，論文提出了一種動態Vision Transformer(DVT)框架，能夠根據每個圖片自動配置合適的token數，實現高效計算。訓練時使用逐漸增多的token數訓練級聯Transformer，測試時從較少的token數開始依次推理，得到置信度足夠的預測即終止推理過程。通過自動調整token數，“簡單”樣本和“困難”樣本的計算消耗將會不一樣，從而顯著提高效率。

另外，論文還設計了基於特徵和基於關係的兩種複用機制，減少冗餘的計算。前者允許下游模型在先前提取的深度特徵上進行訓練，而後者允許利用上游模型中的自注意力關係來學習更準確的注意力圖。

DVT是一個通用框架，可整合到大多數影象識別的Transformer模型中。而且可以通過簡單地調整提前終止標準，線上調整整體計算成本，適用於計算資源動態波動或需要以最小功耗來實現特定效能的情況。從ImageNet和CIFAR的實驗結果來看，在精度相同的情況下，DVT能將T2T-ViT的計算成本降低1.6-3.6倍，而在NVIDIA 2080Ti上的真實推理速度也與理論結果一致。

Dynamic Vision Transformer

Overview

Inference

DVT的推理過程如圖2所示。對於每張測試圖片，先使用少量1D token序列對其進行粗略表示，可通過直接使用分割影象塊或利用如tokens-to-token模組之類的技術來實現，然後通過Vision Transformer對這些token進行快速預測。由於Transformer的計算消耗與token數量成二次增長，所以這個過程很快。最後基於預設的終止標準對預測結果進行快速評估，確定是否足夠可靠。

如果預測未能滿足終止標準，原始輸入影象將被拆分為更多token，再進行更準確、計算成本更高的推理。每個token embedding的維度保持不變，只增加token數量，從而實現更細粒度的表示。此時推理使用的Vision Transformer與上一級具有相同架構，但引數是不同的。根據設計，此階段在某些“困難”測試圖片上權衡計算量以獲得更高的準確性。為了提高效率，新模型可以複用之前學習的特徵和關係。在獲得新的預測結果後，同樣根據終止標準進行判斷，不符合則繼續上述過程，直到結果符合標準或已使用最終的Vision Transformer。

Training

訓練時，需保證DVT中所有級聯Vision Transformer輸出正確的預測結果，其優化目標為：

其中，為訓練集中的一個樣本及其對應的標籤，採用標準的交叉熵損失函式，而表示第個模型輸出的softmax預測概率。

Transformer backbone

DVT是一個通用且靈活的框架，可以嵌入到大多數現有的Vision Transformer模型（如ViT、DeiT和T2T-ViT）之中，提高其效能。

Feature and Relationship Reuse

DVT的一個重要挑戰是如何進行計算的複用。在使用的具有更多token的下游Vision Transformer時，直接忽略之前模型中的計算結果顯然是低效的。雖然上游模型的token數量較少，但也提取了對預測有價值的資訊。因此，論文提出了兩種機制來複用學習到的深度特徵和自注意力關係，僅增加少量的額外計算成本就能顯著提高準確率。

Background

介紹前，先重溫一下Vision Transformer的基本公式。Transformer encoder由交替堆疊的多頭自注意力(MSA)和多層感知器 (MLP)塊組成，每個塊的之前和之後分別添加了層歸一化(LN)和殘差連線。定義表示第層的輸出，其中是樣本的token數，是token的維度。需要注意的是，，對應影象塊和可學習的分類token。假設Transformer共層，則整個模型的計算可表示為：

得到最終的結果後，取其中的分類token通過LN層+全連線層進行最終預測。這裡省略了position embedding的細節，論文沒有對其進行修改。

Feature reuse

DVT中的所有Transformer都具有相同的目標，即提取關鍵特徵進行準確識別。因此，下游模型應該在上游模型計算的深度特徵的基礎上學習才是最高效的，而不是從頭開始提取特徵。為此，論文提出了圖3的特徵複用機制，利用上游Transformer最後輸出的結果來生成下游模型每層的輔助embedding輸入：

由LN+MLP()開頭，對上游模型輸出進行非線性轉換。轉換後將結果reshape到原始影象中的相應位置，然後上取樣並展平來匹配下游模型的token數量。一般情況下，使用較小的以便快速生成。

之後將拼接到下游模型對應層的中間特徵作為預測的先驗知識，也就是將公式3替換為：

與中間特徵拼接，LN 的維度和MLP的第一層從增加到。由於是基於上游輸出生成的，token數少於，它實際上為中的每個token總結了輸入影象的上下文資訊。因此，將命名為上下文embedding。此外，論文發現不復用分類token對效能有提升，因此在公式5中將其填充零。

公式4和5允許下游模型在每層靈活地利用內的資訊，從而最小化最終識別損失，這種特徵重用方式也可以認為隱式地擴大了模型深度。

Relationship reuse

Vision Transformer的關鍵在於自注意力模組能夠整合整個影象的資訊，從而有效地模擬影象中的長距離關係。通常情況下，模型需要在每一層學習一組注意力圖來描述token之間的關係。除了上面提到的特徵複用，論文認為下游模型還可以複用上游模型產生的自注意力圖來進行優化。

定義輸入特徵，自注意力模組先通過線性變換得到query矩陣、key矩陣和value矩陣：

其中，、和為權重矩陣。然後通過一個帶有softmax的縮放點乘矩陣運算得到注意力圖，最後根據注意力圖來計算所有token的值：

其中，是或的點積結果維度，為注意力圖。為了清楚起見，這省略了多頭注意力機制的細節，多頭情況下包含多個注意力圖。

對於關係複用，先將上游模型所有層產生的注意力圖（即）拼接起來：

其中，和分別為上游模型中的toekn數和注意力圖數，通常，是多頭注意力的head數，是層數。

下游的模型同時利用自己的token和來構成注意力圖，也就是將公式7替換為：

其中是一個轉換網路，整合提供的資訊來細化下游注意力圖。的架構如圖5所示，先進行非線性MLP轉換，然後上取樣匹配下游模型的注意力圖大小。

公式9雖然很簡單，但很靈活。有兩個可以魔改的地方：

由於下游模型中的每個自注意力模組可以訪問上游模型的所有淺層和深層的注意力頭，可以嘗試通過可學習的方式來對多層的注意力資訊進行加權整合。
新生成的注意力圖和複用注意力圖直接相加，可以嘗試通過可學習的方式來對兩者加權。

還需要注意的是，不能直接使用常規上取樣操作。如圖5所示，假設需要將()的注意力圖對映上取樣到()的大小。由於每一行對應單個token與其他個token的關係，直接對注意力圖上取樣會引入混亂的資料。因此，需要先將行reshape為，然後再縮放到，最後再展平為向量。

Adaptive Infernece

如前面所述，DVT框架逐漸增加測試樣本的token數量並執行提前終止，“簡單”和“困難”影象可以使用不同的token數來處理，從而提高了整體效率。對於第個模型產生的softmax預測，將的最大項與閾值進行比較。如果，則停止並採用作為輸出。否則，將使用更多token數更多的下游模型繼續預測直到最後一個模型。

閾值需要在驗證集上求解。假設一個計算資源有限的批量資料分類場景，DVT需要在給定的計算預算內識別一組樣本。定義和為資料集上使用閾值時的準確度和計算成本，最優閾值可以通過求解以下優化問題得到：

由於公式10是不可微的，論文使用遺傳演算法解決了這個問題。