UniFormer：無縫集成 Transformer，更高效的時空表征學習框架 | 論文解讀

對圖像和視頻上的表征學習而言，有兩大痛點：局部冗余——視覺數據在局部空間、時間、時空鄰域具有相似性，這種局部性質容易引入大量低效的計算；全局依賴——要實現準確的識別，需要動態地將不同區域中的目標關聯，建模長時依賴。而現在的兩大主流模型 CNN 和 ViT，往往只能解決部分問題。CNN 通過卷積有效降低了局部冗余，但受限的感受野難以建模全局依賴；而 ViT 憑借自注意力可捕獲長距離依賴，但在淺層編碼局部特征十分低效，導致過高冗余。

為能同時解決上述兩大痛點，上海人工智能實驗室聯合商湯科技共同提出一種新的 UniFormer（Unified Transformer）框架， 它能夠將卷積與自注意力的優點通過 Transformer 進行無縫集成。與經典的 Transformer 模塊不同，UniFormer 模塊的相關性聚合在淺層與深層分別配備了局部全局token，能夠同時解決冗余與依賴問題，實現了高效的特征學習。包括 ICLR2022 接受的 video backbone，以及為下游密集預測任務設計的拓展版本，UniFormer 整套框架在各種任務上都能表現出比現有 SOTA 模型更好的性能。

論文標題：UniFormer: Unified Transformer for Efficient Spatiotemporal Representation Learning

論文鏈接：https://arxiv.org/abs/2201.04676

擴展版本：UniFormer: Unifying Convolution and Self-attention for Visual Recognition

代碼鏈接：https://github.com/Sense-X/UniFormer

*UniFormer 性能提升

1、圖像分類：在 Token Labeling 的加持下，僅靠 ImageNet-1K 訓練，39GFLOPs 的 UniFormer-L-384 在 ImageNet 上實現了 86.3% 的 top-1 精度。

2、視頻分類：僅用 ImageNet-1K 預訓練，UniFormer-B 在 Kinetics-400 和 Kinetics-600 上分別取得了 82.9% 和 84.8% 的 top-1 精度（比使用 JFT-300M 預訓練，相近性能的 ViViT 的 GFLOPs 少 16 倍）。在 Something-Something V1 和 V2 上分別取得 60.9%和 71.2% 的 top-1 精度，為同期模型的 SOTA。 

3、密集預測：僅用 ImageNet-1K 預訓練，COCO 目標檢測任務上取得了 53.8 的 box AP 與 46.4 的 mask AP；ADE20K 語義分割任務上取得了 50.8 的 mIoU；COCO 姿態估計任務上取得了 77.4 的 AP。后文將會介紹為下游任務設計的訓練和測試時模型適配。

圖1：圖像分類與視頻分類任務性能比較（上方為 ImageNet 上 224x224 與 384x384 分辨率輸入）

設計靈感

對圖像和視頻上的 representation learning 而言，目前存在兩大痛點：

? Local Redundancy: 視覺數據在局部空間/時間/時空鄰域具有相似性，這種局部性質容易引入大量低效的計算。

? Global Dependency: 要實現準確的識別，需要動態地將不同區域中的目標關聯，建模長時依賴。

現有的兩大主流模型 CNN 和 ViT，往往只關注解決以上部分問題。Convolution 只在局部小鄰域聚合上下文，天然地避免了冗余的全局計算，但受限的感受也難以建模全局依賴。而 self-attention 通過比較全局相似度，自然將長距離目標關聯，但通過如下可視化可以發現，ViT 在淺層編碼局部特征十分低效。

圖2：DeiT 可視化后，可以發現即便經過三層 self-attention，輸出特征仍保留了較多的局部細節。任選一個 token 作為 query，可視化 attention 矩陣可以發現，被關注的 token 集中在 3x3 鄰域中（紅色越深表示關注越多）。

圖3：TimeSformer 可視化后，同樣可以發現即便是經過三層 self-attention，輸出的每一幀特征仍保留了較多的局部細節。任選一個 token 作為 query，可視化 spatial attention 和 temporal attention 矩陣都可以發現，被關注的 token 都只在局部鄰域中（紅色越深表示關注越多）。

無論是 spatial attention 抑或是 temporal attention，在 ViT 的淺層，都僅會傾向于關注 query token 的鄰近 token。但是 attention 矩陣是通過全局 token 相似度計算得到的，這無疑帶來了大量不必要的計算。相較而言，convolution 在提取這些淺層特征時，無論是在效果上還是計算量上都具有顯著的優勢。那么為何不針對網絡不同層特征的差異，設計不同的特征學習算子，從而將 convolution 和 self-attention 有機地結合物盡其用呢？

本論文中提出的UniFormer (Unified transFormer)，旨在以 Transformer 的風格，有機地統一 convolution 和 self-attention，發揮二者的優勢，同時解決 local redundancy 和 global dependency 兩大問題，從而實現高效的特征學習。

模型架構

圖4：模型整體框架，標紅維度僅對 video 輸入作用，對 image 輸入都可視作 1

模型整體框架如上圖所示，借鑒了 CNN 的層次化設計，每層包含多個 Transformer 風格的 UniFormer block。

1、UniFormer block

每個 UniFormer block 主要由三部分組成，動態位置編碼 DPE、多頭關系聚合器 MHRA 及 Transformer 必備的前饋層 FFN，其中最關鍵的為多頭關系聚合器：

2、MHRA

與多頭注意力相似，將關系聚合器設計為多頭風格，每個頭單獨處理一組 channel 的信息。每組的 channel 先通過線性變換生成上下文 token  ，然后在 token affinity   的作用下，對上下文進行有機聚合。

基于前面的可視化觀察，可以認為在網絡的淺層，token affinity 應該僅關注局部鄰域上下文，這與 convolution 的設計不謀而合。因此，將局部關系聚合  設計為可學的參數矩陣：

3、local MHRA

其中  為 anchor token，  為局部鄰域  任一 token，  為可學參數矩陣，(i ? j) 為二者相對位置，表明 token affinity 的值只與相對位置有關。這樣 local UniFormer block 實際上與 MobileNet block 的設計風格相似，都是 PWConv-DWConv-PWConv（見原論文解析），不同的是研究引入了額外的位置編碼以及前饋層，這種特別的結合形式有效地增強了 token 的特征表達。

在網絡的深層，需要對整個特征空間建立長時關系，這與 self-attention 的思想一致，因此，通過比較全局上下文相似度建立 token affinity：

4、global MHRA

其中  為不同的線性變換。先前的 video transformer，往往采用時空分離的注意力機制 ，以減少 video 輸入帶來的過量點積運算，但這種分離的操作無疑割裂了 token 的時空關聯。相反，UniFormer 在網絡的淺層采用 local MHRA 節省了冗余計算量，使得網絡在深層可以輕松使用聯合時空注意力，從而可以得到更具辨別性的 video 特征表達。再者，與以往 ViT 中使用絕對位置編碼不同，這里采用卷積風格的動態位置編碼，使得網絡可以克服 permutation-invariance 的同時，對不同長度的輸入更友好。

流行的 ViT 往往采用絕對或者相對位置編碼 ，但絕對位置編碼在面對更大分辨率的輸入時，需要進行線性插值以及額外的參數微調，而相對位置編碼對 self-attention 的形式進行了修改。為了適配不同分辨率輸入的需要，采用了最近流行的卷積位置編碼設計動態位置編碼:

5、DPE

其中 DWConv 為零填充的的深度可分離卷積。一方面，卷積對任何輸入形式都很友好，也很容易拓展到空間維度統一編碼時空位置信息。另一方面，深度可分離卷積十分輕量，額外的零填充可以幫助每個 token 確定自己的絕對位置。

整體框架

1、圖像分類

表1：模型細節

團隊設計了三種不同規模的模型，每個模型包含 4 層，前兩層使用 local MHRA，后兩層使用 global MHRA。對于 local MHRA，卷積核大小為 5x5，歸一化使用 BN（使用 LN 性能較差）。對于 global MHRA，每個 head 的 channel 數為 64，歸一化使用 LN。動態位置編碼卷積核大小為 3x3，FFN 的拓展倍數為 4。對于特征下采樣，團隊采用非重疊卷積，其中第一次下采樣卷積核大小為 4x4、步長為 4x4，其余三次下采樣卷積核大小為 2x2、步長為 2x2。在每次下采樣卷積之后，額外增加 LN 歸一化。網絡最后接平均池化層與線性分類層，輸出最終預測。當使用 Token Labeling 時，額外加入一個線性分類層以及輔助損失函數。

對于 UniFormer-S，團隊設計了增強版本，每層 block 數量為 [3, 5, 9, 3]，下采樣使用重疊卷積，FLOPs 控制為 4.2G，保證與其他 SOTA 模型可比。

2、視頻分類

對于 video 使用的 3D backbone，通過加載 ImageNet-1K 預訓練的  UniFormer-S 和 UniFormer-B，并進行卷積核展開。具體來說，動態位置編碼和 local MHRA 分別展開為 3x3x3 和 5x5x5 卷積。對于下采樣層，只在第一次下采樣同時壓縮時間和空間維度，而在其余三次下采樣僅壓縮空間維度，即第一次下采樣卷積核大小為 3x4x4、步長為 2x4x4，其余卷積核大小為 1x2x2、步長為 1x2x2。這樣就可以在減少計算量的同時，還能保證模型的高性能。對于 global MHRA，直接繼承相應參數，將時空 token 序列化進行統一處理。

3、密集預測

對于下游密集預測任務，直接使用原始網絡作為 backbone 并不合適。因為這些任務往往輸入大分辨率圖像，比如目標檢測中輸入 1333x800 的圖像，使用 global MHRA 會帶來過多的計算量。以 UniFormer-S 為例，統計了不同分辨率輸入時，不同操作所需的計算量。

圖5：左邊圖片為模型整體計算量與第三/四層中 MatMul 運算所需計算量的結果；右邊圖片為第三層采用不同風格的 MHRA 所需的 MatMul 計算量的結果

從上圖中可以看到，第三層中 MHRA 所需的 MatMul 運算隨著分辨率的增加而急劇上升，在輸入分辨率為 1008x1008 時，甚至占了總運算量 50% 以上，而第四層僅為第三層的 1/28。因此，僅對第三層的 MHRA 進行改進。

受先前工作的啟發，將 global MHRA 應用在限制的窗口內，這樣會把原本  的復雜度降至  ，其中 p 為窗口大小。然而直接應用純 window 化操作，不可避免地會帶來性能下降，為此團隊將window 和 global 操作結合。每個 hybrid 分組中包含 4 個 block ，前 3 個為 window block，最后1 個為 global block。UniFormer-S 和 UniFormer-B 分別包含 2 個和 5 個分組。

圖6：不同任務在訓練和測試時使用特定改進 

如上圖所示，在五種任務上，訓練和測試采用特定改進。對目標檢測任務，由于訓練和測試時輸入分辨率都很大（如 1333x800），在第三層都采用 hybrid block。對姿態估計任務， 輸入分辨率相對較小（如 384x288），在第三層采用原本的 gloabl block。而對于語義分割任務，由于往往在測試時使用幾倍于訓練輸入的分辨率（如 2048x512 vs. 512x512），因此在訓練時，對第三層采用 global block，而在測試時采用 hybrid block，但需要注意測試時 hybrid block 中 window size 需要與訓練時 global block 的感受野一致（如 32x32），感受野不一致會導致急劇性能下降。這種設計可以保證訓練高效的同時，提高測試的性能。

實驗結果

1、圖像分類

團隊用ImageNet-1K 進行圖像分類實驗，采用了 DeiT 的代碼與訓練策略，UniFormer-S/B/L 使用 droppath 比例為 0.1/0.3/0.4，對大模型額外加入 Layer Scale 防止網絡訓練崩潰 。對于 Token Labeling，沿用了其代碼訓練框架與超參。

表2 ：ImageNet-1K

結果如上所示，其中帶*為加入 Token Labaleing 進行訓練，UniFormer-B dagger 為前述 UniFormer-B 的增強版本。可以看到 UniFormer 在不同規模下都取得了 SOTA 性能，UniFormer-L-384 僅需 39G 的 FLOPs，即可取得 86.3% 的 top-1 精度。

2、視頻分類

團隊在 Kinetics-400/600 以及 Something-Something V1/V2 上進行了視頻分類實驗，沿用了 MViT 的代碼和訓練策略。對 Kinetics 采用 dense 采樣方式，加載了 ImageNet-1K 的預訓練模型加速訓練，droppath 比例保持與前述一致。對 Sth-Sth 采用 uniform 采樣方式，加載了 Kinetics 的預訓練，droppath 設置為前述兩倍，并不進行水平翻轉。

表3（左）：Kinetics；表4（右）：Sth-Sth

結果如上圖所示，僅使用 ImageNet-1K 預訓練，結果在 Kinetics 上取得了與使用 JFT-300M 預訓練的 ViViT-L、使用 ImageNet-21K 預訓練的 Swin-B 相近的性能，計算量大幅減小。而在 Sth-Sth 上，UniFormer 大幅高于先前 CNN 和 ViT 的結果，取得了新的 SOTA 結果，在 Sth-Sth V1 上為 61.2%，V2 上為 71.4%。

3、目標檢測與實例分割 

團隊在 COCO2017 上進行了目標檢測和實例分割實驗，沿用了 mmdetection 的代碼框架，配置了 Mask R-CNN 與 Cascade Mask R-CNN 框架，使用 Swin Transformer 的訓練參數，均取得了 SOTA 性能。

表5 ：COCO 目標檢測，Mask R-CNN

表6 ：COCO 目標檢測，Cascade Mask R-CNN

4、語義分割 

團隊在 ADE20K 上進行了語義分割實驗，沿用了 mmsegmentation 的代碼框架，配置了Semantic FPN 與 UperNet 兩種框架，分別使用了 PVT 和 Swin 的訓練參數，均取得了SOTA 性能。

表7、表9 ：ADE20K 語義分割，左：Semantic FPN；右：UperNet

5、人體姿態估計 

團隊在 COCO2017 上進行了姿態估計實驗，沿用了 mmpose 的代碼框架，配置了簡單的 Top-down 框架，使用了 HRFormer 的訓練參數，取得了 SOTA 性能。

表8：COCO 姿態估計，Top-down

6、消融實驗

團隊進行了詳盡的消融實驗，首先在圖像和視頻分類任務上驗證了 backbone 的結構設計。其次，對 video backbone 的預訓練、訓練以及測試策略進行了探究。最后，驗證了下游密集預測任務上改進的有效性。

Model designs for image and video backbones

表10 ：結構設計

? FFN：首先將 local block 替換為 MobileNet block，其中 ReLU 替換為 GeLU，expand ration 設置為 3 保持計算量一致，保留動態編碼以公平比較，可以發現 local block 在圖像 和視頻分類任務上都會明顯優于 MobileNet block。由此可知 transformer 風格，以及其特有的 FFN 確實增強了 token 的特征表達。

? DPE：將動態位置編碼去掉，在圖像和視頻任務上性能均下降，視頻任務上更是掉了 1.7%，由此可知位置編碼有助于更好的時空特征學習。

? Local MHRA size：將 local MHRA 的卷積核大小設置為 3、5、7、9，性能差異并不大，最終采用大小為 5 的卷積核，以取得最好的計算量與準確率的權衡。

? MHRA configuration：由純 local MHRA（LLLL）出發，逐層替換使用 global MHRA。結果可以發現，僅使用 local MHRA 時，計算量很小，但性能下降明顯。逐層替換 global MHRA 后，性能逐漸提升。但全部替換為 global MHRA 后，視頻分類準確率急劇下降， 計算量急劇上升，這主要是因為網絡缺失了提取細節特征的能力，冗余的 attention 在有限的 video 數據下導致了急劇的過擬合。

Pre-training, training and testing for video backbone

表11、12、13、14 ：遷移性能、卷積核展開方式、訓練采樣方式

? Transfer learning：表 11 比較了不同結構的預訓練性能以及小數據集上的遷移學習性能，可以發現，聯合的時空學習方式，不僅在預訓練模型上性能更好，在小數據集上的遷移學習性能提升明顯。而純 local MHRA 以及時空分離的學習方式，遷移小數據訓練未能帶來提升。

? Infalting methods：表 12 中比較了是否對卷積核進行展開，可以發現，展開為 3D 卷積核，在場景相關的數據集 Kinetics-400 上性能接近，但在時序相關的數據集 Sth-Sth V1 上提升明顯，尤其是在強的預訓練的加持下，這表明 3D 卷積核對時空特征的學習能力更強。

? Sampling strides of dense sampling：表 13 中比較了在訓練 Kinetics 時，使用不同間隔采樣的結果。可以發現更稀疏的采樣，在單 clip 測試上效果往往更好，但在多 clip 測試時，間隔 4 幀采樣更好。 

? Sampling methods of Kinetics pre-trained model：由于加載 Kinetics 預訓練模型訓練 Sth-Sth，而 Sth-Sth 采用 uniform 采樣，所以有必要知道預訓練覆蓋更多幀是否能帶來提升。表 14 的結果表明，預訓練的不同采樣方式差別并不大，16x4 采樣在大部分條件下性能都較好。

圖7（左）：不同數據集預訓練。 圖8（右）：不同測試策略。

? Pre-trained dataset scales：圖 7 比較了不同規模數據預訓練的結果，可以發現對于小模型，大數據集預訓練的提升非常明顯，而對于大模型則相差無幾。 

? Testing strategies：圖 8 比較了不同的測試策略，可以發現對于使用 dense 采樣方式訓練的場景相關數據集 Kinetics 而言，多 clip 測試方案較好，且 1x4 綜合性能最優。對于使用 uniform 采樣方式訓練的時序相關數據集 Sth-Sth 而言，多 crop 測試方案較好，且 3x1 綜合性能最好。

Adaption designs for downstream tasks

表15、16、17：下游改進

? object detection：表 15 比較了目標檢測任務上，第三層采用不同類型 block 的結果。盡管在 1x 訓練時，hybrid block 的性能比純 global block 性能略差，但經過 3x 的充分訓練后，hybrid block 的性能已經能和純 global block 持平。

? semantic segmentation：表 16 比較了語義分割任務上，第三層采用不同類型 block 的結果。可以發現更大的窗口，以及 global block 的使用都能明顯提升性能，由于純 global block 計算量較大，團隊采用性能相近的 hybrid block。

? pose estimation：表 17 分別比較了姿態估計任務上，第三層采用不同類型 block 的結果。由于圖像分辨率較小，zero padding 消耗了更多的計算量。

Visualizations

圖9：圖像分類與視頻分類

圖10：目標檢測、語義分割、姿態估計

總結與思考

在過去一年多，研究團隊在視頻模型設計上嘗試了 CNN（CTNet，ICLR2021）、ViT（UniFormer，ICLR2022）及 MLP（MorphMLP，arxiv）三大主流架構，發現：Transformer 風格的模塊+CNN 的層次化架構+convolution 的局部建模+DeiT 強大的訓練策略，保證了模型的下限不會太低。但相比 convolution 以及 linear layer 而言，self-attention 的性價比仍是最高的，同等計算量尤其是小模型下，self-attention 帶來的性能提升明顯，并且對大數據集預訓練的遷移效果更好。

不過傳統 ViT 對不同分辨率輸入并不友好，并且對大分辨率輸入的計算量難以承受，這在 UniFormer 中都盡可能以簡潔的方式解決，DWConv 有限制地引入，也并不會導致過多的顯存開銷與速度延遲，最后在不同的任務上都能取得很好的性能權衡。UniFormer 提供了一個盡可能簡單的框架，也希望后面的研究工作能在這個框架的基礎上，去考慮視頻中的運動信息、時空維度的冗余性、幀間的長時關系建模等等更復雜的問題，實現更大的突破。