Rethinking Vision Transformers for MobileNet Size and Speed
===
###### tags: `Vit` `SuperNet` `EdgeAI`
---
* 論文出處: https://arxiv.org/pdf/2212.08059.pdf
* Code: https://github.com/snap-research/EfficientFormer
---
[TOC]

## 1. Introduction
隨著視覺Transformers（ViTs）在電腦視覺任務中的成功，最近的技術試圖優化ViT的性能和復雜性，以實現佈署在Edge Device上。研究人員提出了多種方法來加速注意力機制，改進低效設計，或結合mobile-friendly的輕量級卷積來形成混合架構。

然而，ViT及其改良版仍然比輕量級的CNNs具有更差的延遲或更多的參數，即使對於多年前的MobileNet也是如此。實際上，延遲和大小對於資源受限硬件上的高效部署都至關重要。在論文中研究了一個中心問題，ViT模型是否可以像MobileNet一樣快速運行並保持類似的大小？

論文重新審視了ViT的設計選擇，並提出了一種具有低延遲和高參數效率的改進型超網。論文進一步引入了一種fine-grained architecture search algorithm，該策略可以通過同時優化延遲和參數量來找到有效的架構。所提出的模型EfficientFormerV2在ImageNet-1K上實現了比MobileNetV2和MobileNetV1高約4%的top-1精度，具有相似的延遲和參數。論文證明，適當設計和優化的ViT可以以MobileNet級別的大小和速度實現高性能。


總結來說，本文的貢獻如下：

* 論文提供了一項全面的研究，以驗證和改進mobile-friendly設計選擇，這是獲得超高效視覺主幹的實用指南；
* 論文提出了一種搜索算法，該算法同時優化了模型大小和速度，實現了卓越的Pareto最優；
* EfficientFormerV2系列模型實現了超快的推理和極小的模型大小，大大超過了以前的技術，並且可以在各種下游任務中充當強大的骨幹。
## 2. Related Work
論文中提出目前針對ViT的改進研究有兩個主要卻點。
1. 首先許多優化並不是和佈署在Edge Device上，例如:通過歸一化self-attention在二次計算的複雜度可以降低為線性，但在資源受限的設備上甚至不支持複雜的reshaping和索引操作。
2. 儘管最近的混合設計和網路搜索方法展示了具有強大性的高效ViT，但他們主要優化一維的Pareto cure，而在其他維度上卻比較弱，例如:MobileViT的訓練參數少但其在設備上的執行速度卻遠慢於清化後的CNN模型。
## 3. EfficientFormweV2
EfficientFormerV2的模型設計沿用EfficinetFormer架構進行改良，，下圖為整個網路架構：


採用了四階段分層設計，其獲得的特徵尺寸為輸入分辨率的 **1/4，1/8，1/16，1/32**。與其前身類似，EfficientFormerV2從一個小的內核卷積stem開始嵌入輸入圖像，而不是使用非重疊patch的低效嵌入

前兩個階段捕獲高分辨率的局部信息；因此論文只使用統一的FFN（FFN，圖2（b））：

在最後兩個階段，使用局部FFN和全局MHSA塊。因此式4，global blocks定義為：

作者針對原本模型的Stage 1的Local module以及Stage 4 的MHSA進行調整，此外也針對高解析度輸入的圖形進行模型設計的調整，並以EffocorntFormer-L1作為整個實驗的baseline，下面圖表是作者統計的實驗改進結果：

### 3.1 Token Mixers vs. Feed Forward Network
通常來說，結合局部信息可以有效提高性能，同時使 ViTs 對明確的位置嵌入缺失表現得更加穩健。 PoolFormer 和 EfficientFormer 中都使用了 3×3 的平均池化層(如圖 2(a)所示)作為局部的Token Mixers。採用相同卷積核大小的深度可分離卷積(DWCONV)替換，不僅不會帶來延遲開銷，同時性能也能提高 **0.6%**，參數量僅微漲 0.02M。此外，同 NASVit，作者也在 ViTs 的前饋網絡(FFN)中注入了局部信息建模層，這也有益於提高性能。

這裡，作者直接將原來的 Pooling 層刪掉了(下採樣越大，理論感受野越大)，而是直接替換成 BottleNeck 的形式，先用 1×1 卷積降維壓縮，再嵌入 3×3 的深度可分離卷積提取局部信息，最後再通過 1×1 的捲積升維。這樣做的一個好處是，這種修改有利於後續直接用超參搜索技術搜索出具體模塊數量的網路深度，以便在網路的後期階段中提取局部和全局信息。
### 3.2 Search Space Refinement
通過調整網路的深度(即每個階段中的層數)和寬度(即通道數)，可以發現，更深和更窄的網絡可以帶來：

* 更好的準確性(**0.2%** 的性能提升)
* 更少的參數(**0.13M** 的參數壓縮)
* 更低的延遲(**0.1ms** 的推理加速)

同時將這個修改的網絡設置為後續 NAS 搜索的 Baseline.
### 3.3 MHSA Improvements
一般的 Transformer 模塊都會包含兩個組件，即Multihead self-attention (MHSA) 和全連接層 FFN. 作者隨後便研究瞭如何在不增加模型大小和延遲的情況下提高注意module性能的技術。

首先，通過 3×3 的捲積將局部信息融入到 Value 矩陣中，這一步跟 NASVit 和 Inception transformer 一樣。

其次，在 Head 維度之間添加 FC 層(就圖中標識的Talking Head)，增強不同頭之間的信息交互。

通過這些修改，可以發現性能進一步提高了，與baseline模型相比，準確度達到了 80.8%，同時參數和延遲也基本保持一致。
### 3.4 Attention on Higher Resolution
加入self-attention機制普遍都是能夠提升性能的。然而，將它應用於高解析率特徵會降低行動設備的效能，因為它具有與解析度成平方關係的時間複雜度。因此，作者僅在最後的 1/32 的空間分辨率下使用，而對於倒數的第二階段即 4 倍下採樣提出了另外一種 MHSA，這有助於將準確率提高了 0.9% 同時加快推理效率。

先前的解決方案，例如 Cswin transformer 和 Swin transformer 都是採用基於滑動窗口的方式去壓縮特徵維度，又或者像 Next-vit 一樣直接將 Keys 和 Values 矩陣進行下採樣壓縮，這些方法對於行動裝置部署並不是一個最佳的選擇，以 Swin transformer 為例，它在每個 stage 都需要進行複雜的窗口劃分和重新排序，所以這種基於 windows 的注意力是很難在行動裝置上加速優化的。而對於 Next-vit 來說表面上看雖然進行了壓縮，但整個 Key 和 Value 矩陣依舊需要全分辨率查詢矩陣(Query)來保持注意力矩陣乘法後的輸出分辨率：

$$
Out_{[B,H,N,C]} = (Q_{[B,H,N,C]}\cdot K_{[B,H,C,{\frac{N}{2}}]}^T ) \cdot V_{[B,H,{\frac{N}{2}},C]}
$$

本文的解決方案可以參考圖(d)和(e)，整體思路是採用一個stride的注意力，實現上就是將所有的 QKV 均下採樣到固定的空間分辨率(這裡是 1/32)，並將注意力的輸出複原到原始分辨率以餵入下一層，有點類似於把bottleneck layer的思路又搬過來套。根據測試，使用這種方法可以將模型的延遲時間從 3.5ms 降低到 1.5ms，同時準確率僅略有下降（81.5% vs. 81.7%）。
### 3.5 Attention Downsampling
以往的下採樣方式大都是採用帶步長的捲積或者池化層直接進行的。不過最近也有一部分工作在探討 Transformer 模塊內部的下採樣方式，如 LeViT 和 UniNet 提出通過注意力機制將特徵分辨率減半，從而更好的整合全局上下文來進感知下採樣。具體的做法也就是將 Query 中的 Token 數量減半，從而對注意力模塊的輸出進行壓縮：
$$
Out_{[B,H,{\frac{N}{2}},C]} = (Q_{[B,H,{\frac{N}{2}},C]}\cdot K_{[B,H,C,N]}^T ) \cdot V_{[B,H,N,C]}
$$

在這裡不經意間有個疑問，Token 數量減少多少才是合適？況且，如果直接對所有的Query矩陣進行降採樣的話，這對於較前的 stage 的特徵提取是不利的，因為網絡的淺層更多的是提取諸如紋理、顏色、邊緣等 low-level 的信息，因此從經驗上來看是需要保持更高分辨率的。

作者的方法是提出一種結合局部和全局上下文融合的組合策略，如上圖(f)所示。為了得到下採樣的Query，採用池化層作為靜態局部下採樣，而 3×3 DWCONV 則作為可學習的局部下採樣，並將結果拼接起來並投影到查詢矩陣中。此外，注意力下採樣模塊殘差連接到一個帶步長的捲積以形成局部-全局方式，類似於downsampling bottlenecks 或inverted bottlenecks。

## 4.Jointly Optimizing Model Size and Speed
作者提到了 EfficientFormer 的兩個主要缺陷。首先，搜索過程僅受到速度的限制，導致最終模型參數冗餘。其次，它僅搜索每個階段的深度（塊數 ${N_j}$）和寬度 $C_j$，這是一種coarse-grained的方式。實際上，網絡的大多數計算和參數都在 FFN 中，並且參數和計算複雜度與其擴展率 $E_{i,j}$ 線性相關。因此，搜索 $E_{i,j}$ 可以更fine-grained的search space，其中計算和參數可以在每個階段內靈活且非均勻分佈，這也是最近 ViT NAS 中缺少的特性。作者提出了一種搜索算法，該算法使得每個Block配置更加靈活，並在大小和速度之間互相約束，並找到最適合行動設備的 vision backbones。

### 4.1 Search Objective
作者介紹了指導其聯合搜索算法的衡量標準。由於在評估mobile-friendly models時網絡的大小和延遲都很重要，因此我們考慮使用一種通用且公平的衡量標準，更好地了解網絡在移動設備上的性能。為了不失一般性，我們定義了一個 Mobile Efficiency Score（MES）：
$$
MES = Score\cdot \prod_{i}^{}(\frac{M_i}{U_i}) ^{-\alpha_i}
$$
MES 衡量了網路在行動設備上的效率，它既考慮了網絡的速度，也考慮了網絡的大小。作者認為，通過使用 MES 這種度量標準，可以更好地找到最適合行動設備的網絡。

### 4.2 Search Space and SuperNet
作者討論了搜索空間的構成。搜索空間包括：(i)網絡的深度，用每個階段的塊數 Nj 測量，(ii)網絡的寬度，即每個階段的通道維度 Cj，以及(iii)每個 FFN 的擴展比 Ei,j。 MHSA 的數量可以在深度搜索期間輕鬆確定，它控制超網中塊的保留或刪除。因此，作者在超網的最後兩個階段中將每個塊設置為 MHSA，然後跟隨 FFN，並通過深度搜索獲得具有所需全局 MHSA 數量的子網絡。

作者進一步討論了超網的構建方式。作者使用了slimmable network來構建超網，該網絡具有彈性深度和寬度，從而可以執行基於評估的純搜索算法。彈性深度可以通過隨機丟棄路徑增強（stochastic drop path augmentation）自然實現。對於寬度和擴展比，我們遵循 Yu 等人的方法構建具有共享權值但獨立歸一化層的可切換層，使得相應層可以從預定義集合中的不同通道數執行，即 16 或 32 的倍數。

### 4.3 Search Algorithm
作者介紹了在已確定搜索目標、搜索空間和超網的情況下使用的搜索算法。由於超網可以在彈性深度和可切換寬度下執行，因此可以通過分析每次縮減操作的效率收益和精度下降來搜索具有最佳帕累托曲線（Pareto curve）的子網絡。作者定義了操作池（action pool）如下：1.每個塊的寬度縮減，2. 每個塊的深度減少，3. 刪除操作池中的塊。通過這些操作，作者可以找到一個子網絡，其具有最優的帕累托曲線，即在保持精度的同時盡可能降低模型的大小和推理時間。


## 5. Experiments
### 5.1 Evaluation on Single Metric
EfficientFormerV2系列在單個指標（即參數數量或耗時）上實現了最先進的結果。在模型大小方面，EfficientFormerV2-S0比EdgeViT XXS的精度高出1.3%，參數甚至少了0.6M，比MobileNetV2×1.0的精度高3.9%，參數數量相似。對於大模型，EfficientFormerV2-L模型實現了與最近的EfficientFormer-L7相同的精度，同時小3.1倍。在速度方面，EfficientFormerV2-S2具有可比或更低的延遲，分別以0.8%、0.6%和2.4%的top-1精度優於UniNet-B1、EdgeViT-S和EfficientFormer-L1。


### 5.2 Jointly Optimized Size and Speed
當考慮模型大小和速度時，論文證明了EfficientFormerV2的優越性能。在這裡使用MES作為評估移動效率的更實用的指標，而不是僅使用大小或耗時。 EfficientFormerV2-S1分別以4.2%、4.6%和1.5%的優勢超過MobileViT-XS、EdgeViT-XXS和EdgeViT-XS。憑藉1.8倍更高的MES，EfficientFormerV2-L比MobileFormer-508M高出4.0%的top-1精度。 MES與準確度的可視化如圖3所示。


### 5.3 Downstream Tasks
目標檢測和實例分割：選擇Mask-RCNN作為自己先，實驗結果如表3所示。


### 5.4 Ablation on Search Algorithm
論文將所提出的搜索算法與EfficientFormer的普通算法進行了比較。如表4所示，論文的搜索算法獲得的模型參數和耗時與EfficientFormer相似，但精度更高，證明了細粒度搜索和耗時和大小聯合優化的有效性。


## 6. Discussion and Conclusion
這裡研究了混合視覺骨幹網絡並驗證了適用於移動設備的設計選擇。作者進一步提出了一種針對大小和速度的fine-grained search，並獲得了 EfficientFormerV2 模型系列，該模型既輕量又在推理速度上極快。由於本文簡單地關注了大小和速度，因此未來的一個方向是將聯合優化方法應用於探索其他關鍵指標（如內存佔用和二氧化碳排放）的後續研究中。