導讀

本文提出了一種通用的 CNN 和 Transformer 混合的視覺基礎模型，移動設備和 ImageNet 數據集上的精度相同的前提下，FastViT 比 CMT 快3.5倍，比 EfficientNet 快4.9倍，比 ConvNeXt 快1.9倍

太長不看版

本文是 MobileOne 原班人馬打造，可以看做是 MobileOne 的方法在 Transformer 上的一個改進型的應用。作者取名 FastViT，是一種 CNN，Transformer 混合架構的低延時模型。作者引入了一種新的 token mixer，叫做 RepMixer，它使用結構重新參數化技術，通過刪除網絡中的 Shortcut 來降低內存訪問成本。

進一步使用大核卷積使得 FastViT 精度得到提升，而且不怎么影響延時。在移動設備和 ImageNet 數據集上的精度相同的前提下，FastViT 比 CMT 快3.5倍，比 EfficientNet 快4.9倍，比 ConvNeXt 快1.9倍。在類似的延遲下，FastViT 在 ImageNet 上獲得的 Top-1 準確率比 MobileOne 高 4.2%，是一種極具競爭力的混合架構模型。

1 FastViT：快速卷積 Transformer 的混合視覺架構

論文名稱：FastViT: A Fast Hybrid Vision Transformer using Structural Reparameterization

在CVer微信公眾號后臺回復：FastViT，可以下載本論文pdf、代碼

論文：https://arxiv.org/pdf/2303.14189

代碼：https://github.com/apple/ml-fastvit

1.1 背景和動機

本文的目標是做一個卷積，Attention 的低延時混合架構，因為這種架構有效地結合了 CNN 和 Transformer 的優勢，在多個視覺任務上有競爭力。本文的目標是建立一個模型，實現 SOTA 的精度-延時 Trade-off。

本文的出發點是最近的像 CMT[1]，LIT[2]等 CNN 和 Transformer 混合架構的模型都遵循 MetaFormer[3] 的架構，它由帶有 skip-connection 的 token mixer 和帶有 skip-connection 的前饋網絡 (Feed Forward Network) 組成。由于增加了內存訪問成本 (memory access cost)，這些跳過連接在延遲方面占了很大的開銷。為了解決這個延遲開銷，本文提出 RepMixer，這是一個完全可以重參數化的令牌混合器，它的特點1是使用結構重參數化來刪除 skip-connection。

RepMixer 的特點2是在訓練期間為主要的層添加一些過參數化的額外的分支，以在訓練時提升模型的精度，在推理時全部消除。RepMixer 的特點3是在網絡中使用了大核卷積在前幾個階段替換掉 Self-Attention。具體是在前饋網絡 (FFN) 層和 Patch Embedding 層中加入了大核卷積。這些更改對模型的總體延遲影響很小，同時提高了性能。

對于性能這塊作者在 iPhone 12 Pro 設備和 NVIDIA RTX-2080Ti desktop GPU 上進行了詳盡的分析，實驗結果如下圖1所示。可以看到在兩種設備上，FastViT 都實現了最佳的精度-延時的權衡。

圖1：iPhone 12 Pro 設備和 NVIDIA RTX-2080Ti desktop GPU 上的精度-延時比較

1.2 FastViT 模型架構

FastViT 整體架構如下圖2所示。

Stage 的內部架構

FastViT 采用了4個 stage 的架構，每個 stage 相對于前一個的分辨率減半，通道數加倍。前3個 stage 的內部架構是一樣的，都是訓練的時候采用下式：

推理的時候采用結構重參數化得到下式：

第4個 stage 的內部架構如圖2 (a) 所示，采用 Attention 來作為 token mixer，可能是為了性能考慮，寧愿不采用結構重參數化，犧牲延時成本，以換取更好的性能。

值得注意的是，每個 Stage 中的 FFN 使用的并不是傳統的 FFN 架構，而是如圖2 (c) 所示的，帶有大核 7×7 卷積的 ConvFFN 架構。

圖2：FastViT 模型架構

Stem 的結構

Stem 是整個模型的起點，如圖2 (b) 所示，FastViT 的 Stem 在推理時的結構是 3×3 卷積 + 3×3 Depth-wise 卷積 + 1×1 卷積。在訓練時分別加上 1×1 分支或者 Identity 分支做結構重參數化。

Patch Embedding 的架構

Patch Embedding 是模型在 Stage 之間過渡的部分，FastViT 的 Patch Embedding 如圖2 (a) 所示，在推理時的結構是 7×7 大 Kernel 的 Depth-wise 卷積 + 1×1 卷積。在訓練時分別加上 3×3 分支做結構重參數化。

位置編碼

位置編碼使用條件位置編碼，它是動態生成的，并以輸入 token 的局部鄰域為條件。這些編碼是由 depth-wise 運算符生成的，并添加到 Patch Embedding 中。

1.3 RepMixer 的延時優勢

如下圖3所示，作者對比了 RepMixer 和高效的 Pooling 操作的延時情況。架構使用的是 MetaFormer S12，大概有 1.8 GFLOPs。作者在 iPhone 12 Pro 移動設備上為從 224×224 到 1024×1024 的各種輸入分辨率的模型計時。可以看到 RepMixer 明顯優于 Pooling，尤其是在較高分辨率的時候。在分辨率為 384×384 時，使用 RepMixer 可以降低 25.1% 的延遲，而在分辨率為 1024×1024 時，使用 RepMixer 可以降低 43.9% 的延遲。

圖3：RepMixer 的延時優勢

1.4 FastViT 的大核卷積

RepMixer 的感受野是局部的。我們知道 Self-Attention 操作的感受野是全局的，但是 Self-Attention 操作計算量昂貴，因此之前有工作說使用大核卷積可以在計算量得到控制的情況下有效增加感受野的大小。FastViT 在兩個位置引入了大核卷積，分別是 Patch Embedding 層和 FFN。對比實驗的結果如下圖4所示。將 V5 與 V3 進行比較，模型大小增加了 11.2%，延遲增加了 2.3 倍，而 Top-1 精度的增益相對較小，只有 0.4%，說明使用大核卷積來替換 Self-Attention 是一種高效，節約延時的方式。V2 比 V4 大 20%，延時比 V4 高 7.1%，同時在 ImageNet 上獲得相似的 Top-1 精度。

圖4：大核卷積的消融實驗

隨著感受野的增加，大核卷積也有助于提高模型的魯棒性。FastViT 各種模型的超參數配置如下圖5所示。

圖5：FastViT 的超參數配置

1.5 實驗結果

ImageNet-1K 圖像分類實驗結果

如下圖6所示是 ImageNet-1K 圖像分類實驗結果。對于 iPhone 設備延時的測量，作者使用 Core ML Tools (v6.0) 導出模型，并在帶有 iOS 16 的 iPhone12 Pro Max 上運行，并將所有模型的 Batch Size 大小設置為1。對于 GPU延時的測量，作者把模型導出為 TensorRT (v8.0.1.6) 格式，并在 NVIDIA RTX-2080Ti 上運行，Batch Size 大小為8，報告100次運行的中位數。

與 SOTA 模型的性能比較如下圖6所示。本文的 FastViT 實現了最佳的精度-延時均衡，比如 FastViT-S12 在 iPhone 12 Pro 上比 MobileOne-S4 快 26.3%，GPU 上快 26.9%。在 83.9% 的 Top-1 精度下，FastViT-MA36 比 iPhone 12 Pro 上優化的 ConvNeXt-B 模型快 1.9倍, GPU上快2.0倍。

圖6：ImageNet-1K 圖像分類實驗結果

知識蒸餾實驗結果

如下圖7所示是 FastViT 作為學生模型的知識蒸餾實驗結果。作者遵循 DeiT 中的實驗設置，RegNet16GF 作為教師模型，使用 Hard Distillation，其中教師的輸出設置為 true label，一共訓練300個 Epochs。FastViT 優于最近最先進的模型 EfficientFormer。FastViT-SA24 的性能與 EfficientFormer-L7 相似，但參數少3.8倍，FLOPs 少2.7倍，延遲低2.7倍。

圖7：知識蒸餾實驗結果

目標檢測和語義分割實驗結果

對于語義分割，作者在 ADE20k 上驗證了模型的性能語義分割模型頭使用的是 Semantic FPN，所有的模型都是用預先訓練好的對應圖像分類模型的權重進行初始化。在 512×512 的設置上估計 FLOPs 和延遲。由于輸入圖像的分辨率較高，在表9和表10中，GPU 延遲在測量時使用了大小為2的 Batch Size。在圖8中，作者將 FastViT 與最近的工作進行了比較。FastViT-MA36 的 mIoU 比 PoolFormer-M36 高 5.2%，但是 PoolFormer 具有更高的 FLOPs、參數量和延遲。

圖8：語義分割實驗結果

目標檢測和實例分割實驗實驗 MS-COCO 數據集，實驗結果如下圖9所示。所有模型都使用 Mask-RCNN 目標檢測和實例分割頭按照 1x schedule 進行訓練。所有的模型都是用預先訓練好的對應圖像分類模型的權重進行初始化。結果顯示出 FastViT 在多種延遲機制下實現了最先進的性能。FastViT-MA36 模型的性能與 CMT-S 相似，但在桌面GPU 和移動設備上分別快2.4倍和4.3倍。

圖9：目標檢測和實例分割實驗結果

總結

本文提出了一種通用的 CNN 和 Transformer 混合的視覺基礎模型，是由 MobileOne 原班人馬打造，可以看做是 MobileOne 的方法在 Transformer 上的一個改進型的應用。作者引入了一種新的 token mixer，叫做 RepMixer，它使用結構重新參數化技術，通過刪除網絡中的 Shortcut 來降低內存訪問成本，尤其是在較高分辨率時。作者還提出了進一步的架構更改，以提高 ImageNet 分類任務和其他下游任務的性能。在移動設備和 ImageNet 數據集上的精度相同的前提下，FastViT 比 CMT 快3.5倍，比 EfficientNet 快4.9倍，比 ConvNeXt 快1.9倍。在類似的延遲下，FastViT 在 ImageNet 上獲得的 Top-1 準確率比 MobileOne 高 4.2%，是一種極具競爭力的混合架構模型。