輕量級Visual Transformer模型——LeViT(ICCV2021)

LeViT是FAIR團隊發表在ICCV2021上的成果，是輕量級ViT模型中的標桿，文章對ViT中多個部件進行的改進，如加速策略等，對很多工程化鋪設ViT系列模型都是很有借鑒意義的，
按說，近期出現的優質模型非常多，各種沖擊SOTA的，詳情可戳我整理的小綜述《盤點2021-2022年出現的CV神經網路模型》，但我為何會單獨對LeViT拿出來進行詳細剖析呢？原因很簡單：LeViT非常工程實用，是一款足夠優秀的輕量級視覺transformer模型，市面上很多輕量級模型都進入了一個誤區：大家都在比拼FLOPs數和引數數量，卻忽略了工程上最直觀的評價標準Inference Speed，
首先要知道一個前提就是，無論是FLOPs還是parameters數量，都與inference speed并不成線性關系，
先看一組資料：

即使EfficientNet B0跟LeViT-128有著接近的FLOPs和params，甚至前者這兩項資料更優秀，但推理速度后者可以在GPU上快2倍，CPU上快3倍，ARM CPU上快樂近10倍，看過這個資料，你是否對LeViT有了興趣呢？

按照慣例：
論文標題(鏈接)：LeViT: A Vision Transformer in ConvNet’s Clothing for Faster Inference
Github：https://github.com/facebookresearch/LeViT

LeViT

LeViT的結構如上圖所示，很像一個火箭形狀，這也是LeViT的特點之一：快速縮小特征圖，這里的關鍵點在“特征圖”，對，LeViT的思想里淡化了transformer中“token”的概念，引入了CNN中activation map（特征圖）的概念，所以，大致按對待CNN的思路來對待LeViT會更有助于理解這個模型，比如最后加的那個average pooling，

圖1咱們從下看到上，

1. 用CNN代替“劃片”
這里用到了4層卷積層，每層的步長都是2，這樣不需要使用pooling也能完成特征圖下縮，在這一步，下縮幅度非常sharp，用4層就縮小了16倍的邊長，而ResNet-18縮小到這個尺寸用了10層，如此sharp的尺寸下縮，使得送到self-attention模塊的size就非常小了，有助于加速模型，
3x224x224 -> 256x14x14
輸出為256x14x14的特征圖，咱們用attention的角度理解：256個通道，每個通道14x14個token，咱們要知道self-attention的計算次數跟token數呈二次線性相關（知識點），所以token數太大對自注意力計算非常unfriendly，所以，輕量級ViT的思路肯定是盡量減小輸入到self-attention模塊之前的token數的，
用Conv layers代替無交集劃片有幾個優勢：1. 可以進行初步的特征提取；2. 保留相對位置資訊，可以不用在這個部分進行位置編碼，

2. 用BN代替LN
transformer里最流行的norm方法就是LN，在很多論文的實驗中都表明用LN比BN可以帶來更多的準確率提升，而LeViT作者認為，LN比BN在準確率提升的作用上很有限，但LN比BN在計算上慢一些，
文章里對BN為何更快的解釋如下：在inference的時候，BN計算可以融合到卷積里，幾乎不會額外帶來計算負擔，（這點我附議，在TensorRT實驗的時候可以看到BN的引數都固化下來了，而且前段時間的RepVGG也有融合BN到CNN中的操作，）
原文描述如下：

The batch normalization can be merged with the preceding convolution for inference, which is a runtime advantage over layer normalization (for example, on EfficientNet B0, this fusion speeds up inference on GPU by a factor 2).

3. 多解析度金字塔
這個trick就描述了LeViT整個火箭型的結構，每過一個block，activation maps的尺寸都會減小，從開始的14x14到后面的4x4，

4. 具有收縮特征圖尺寸作用的self-attention模塊

如Figure 5所示，咱們只需要關注輸入輸出的size，左邊是正常的自注意力模塊，右邊是帶特征圖下縮的自注意力模塊，通過在Q的計算里引入下采樣來實作，

5. 采用維度更小的key
在原始ViT里面，K、Q、V都有一樣的維度，實際上V可以保持大一點的維度，我們可以認為V承載了token的主要資訊，而key可以看做是更高層的抽象，所以采用更小的key維度也并不影響模型效果，另外，我們看self-attention計算公式，會發現 K × Q K \times Q K×Q才是self-attention計算的大頭，如果shrink一下key的維度，那自注意力模塊的計算效率將會大大提升，

6. 采用注入attention的可學習位置編碼
這就是LeViT一個很優秀的創新了，傳統的ViT位置編碼只在劃片時候，以人工定義的位置矩陣來進行位置編碼，但是LeViT作者認為，token的相對位置不僅在最開始很重要，在每一層attention模塊都很重要，所以，采用一種注入self-attention計算的位置編碼，咱們看圖：

就是這個位置，用Attention Bias來代替Position Embedding（劃重點），
在看一眼公式：

就是在QK乘法之后加上一個bias常數，這個bias是對稱的、可學習的，

7. 蒸餾學習
蒸餾學習是指采用一個teacher模型來訓練student模型，最開始將distillation思想用到ViT的模型是DeiT，而LeViT就是DeiT的改進版，

transformer引入CV領域之后有個很大的缺點遲遲沒有被解決，那就是transformer不夠Data Efficient，相比CNN，ViT對訓練資料更加依賴，要達到CNN的SOTA，ViT需要用更多的資料以及需要更多的訓練周期，

咱們看上圖，DeiT的收斂比ResNet要慢很多，而且CNN比例越高，模型收斂越快，
這個原因很好解釋：（下面是本人畫的圖）

卷積的優勢是自帶對neighbors的歸納偏執，所以收斂更快，但缺乏global性；
self-attention的優勢是global性，但本身處理單元時，在初始狀態下會對所有token“一視同仁”，如上圖所示，self-attention比Conv天生有更大的解空間，

于是，引入蒸餾訓練就很有必要了，
蒸餾訓練是找個well-trained的CNN當做teacher model，然后用LeViT一方面跟Ground Truth做CE loss，一方面跟CNN的輸出distribution做一個CE loss，這樣可以用CNN來引導LeViT來加快收斂，這樣就可以做到data-efficient了，
LeViT用到的teacher model是RegNetY-16G，
咱們看一下loss function就一目了然了：

上面截圖來自DeiT的論文，蒸餾loss分為軟蒸餾和硬蒸餾兩種，通常用的都是硬蒸餾，