三星提出XFormer | 超越MobileViT、DeiT、MobileNet等模型

ViT 的最新進展在視覺識別任務中取得了出色的表現。卷積神經網絡 (CNN) 利用空間歸納偏差來學習視覺表示，但這些網絡是空間局部的。ViTs 可以通過其self-attention機制學習全局表示，但它們通常是heavy-weight的，不適合移動設備。

在本文中提出了Cross Feature Attention(XFA) 以降低 Transformer 的計算成本，并結合高效的mobile CNNs形成一種新穎的高效輕量級 CNN-ViT 混合模型 XFormer，可作為通用主干學習全局和局部表示。

實驗結果表明，XFormer 在不同的任務和數據集上優于眾多基于 CNN 和 ViT 的模型。在 ImageNet-1K 數據集上，XFormer 使用 550 萬個參數實現了 78.5% 的 top-1 準確率，在相似數量的參數下，比 EfficientNet-B0（基于 CNN）和 DeiT（基于 ViT）的準確率分別提高了 2.2% 和 6.3%。

XFormer的模型在轉移到目標檢測和語義分割任務時也表現良好。在 MS COCO 數據集上，XFormer 在 YOLOv3 框架中超過 MobileNetV2 10.5 AP（22.7 → 33.2 AP），只有 6.3M 參數和 3.8G FLOPs。在 Cityscapes 數據集上，只有一個簡單的 all-MLP ****，XFormer 實現了 78.5 的 mIoU 和 15.3 的 FPS，超過了最先進的輕量級分割網絡。

一個標準的 ViT 模型首先使用patch size h×w 將輸入  reshape為一系列flattened patches ，其中  和  表示token數。然后將  投影到固定的 D 維空間  并使用一堆transformer blocks來學習inter-patch表示。由于忽略了空間歸納偏差，ViT 通常需要更多參數來學習視覺信息。此外，transformers 中 self-attention 的昂貴計算導致優化此類模型的瓶頸。

在本節提出了 XFormer，這是高效、輕量級的 CNN-ViT 框架，以解決 ViT 中的上述問題。首先介紹了一種提高自注意力效率的新方法，然后說明了新的 CNN-ViT 混合模型的架構設計。

Transformer 的主要計算瓶頸之一在于 self-attention 。在最初的 self-attention 過程中， 首先用于通過線性投影生成query Q、key K 和value V。它們都具有相同的維度(N × D)，其中 N 是圖像token數，每個維度為 D。然后計算注意力分數為：

其中 σ 是 softmax 操作， 是head維度。計算注意力分數的計算復雜度為 。 self-attention 的二次復雜性導致了巨大的計算瓶頸，這使得 ViT 模型難以在移動設備上按比例縮小。

為了解決 self-attention 中的二次復雜性問題，作者提出了一種新的注意力模塊結構，稱為Cross Feature Attention（XFA）。在之前的工作之后，首先沿特征維度 D 對query Q 和key K 應用 L2 歸一化：

 在原始 self-attention 中的直覺是定位應該關注的重要圖像塊。

但是直接計算會導致不必要的冗余和計算開銷。相反，作者為 K 構建了2個中間分數：查詢上下文分數  和查詢特征分數 。使用2個卷積核矩陣  和  沿token維度 N 計算 ，沿特征維度 D 計算 。借助卷積濾波器中間分數向量可以表示計算注意力圖的更緊湊的表示，同時也降低了計算成本。 和  表示為：

最后，將Cross Feature Attention(XFA)定義為：

其中 λ 是一個溫度參數，用于動態調整不同transformer層中的比例因子，從而提高訓練穩定性。注意到歸一化將注意力值限制在一定范圍內，因此放棄了冗余且昂貴的 softmax 操作。與原始的具有二次復雜度的 self-attention 不同，XFA模塊將計算成本從  降低到 .

本文提出的 XFA 模塊和原始注意力之間的主要區別是：

1. XFA 通過構建中間查詢上下文和特征分數，沿特征維度D計算注意力圖，大大降低了計算成本；

2. XFA 使用可學習的溫度縮放參數來調整歸一化，并且不受 softmax 操作的影響。

本文方法為二次復雜度問題提供了解決方案，并且對于資源受限的設備更有效且對移動設備更友好。

MobileNetV2 首先引入了inverted residual和linear bottleneck以構建更高效的層結構。MobileNetV3 隨后添加了squeeze excitation (SE) 模塊以處理更大的表示特征。最近的工作證明了通過在 ViT 的早期階段結合卷積層來提高 ViT 性能的合法性。受這種直覺的啟發，繼續為輕量級模型探索這種 CNN-ViT 混合設計。MobileNetV3 塊內的操作可以表述為：

其中  是前一層的輸入特征， 是 MV3 Block 的輸出特征，SE 是squeeze excitation模塊。 表示depth-wise卷積操作， 表示point-wise卷積操作。

利用提出的Cross Feature Attention模塊介紹了 XF Block，一個精心設計的輕量級 transformer module。XF Block內的操作可以表述為：

其中  是前一層的輸入特征， 是 XFA 模塊的特征， 是 XF Block 的輸出特征。LN 表示層歸一化操作，MLP 為全連接層。

對于較大的模型（ViT-Large），基于 ViT 的模型通常采用 8×8、16×16 甚至 32×32 的Patch Size。具有較大Patch Size的優點之一在于，對于分類等圖像級任務，ViT 可以有效地提取圖像塊信息，而不會增加過多的計算開銷。

最近的工作表明，當遷移到語義和實例分割等下游任務時，更小的Patch Size更受青睞，因為它可以增強 Transformer 學習更好的像素級信息的能力，這通常會帶來更好的性能。

此外，隨著Patch Size的減小，token數 N 會大得多。線性復雜度 XFA 模塊可以避免潛在的計算瓶頸。在網絡設計中，每個 XF 塊的Patch Size設置為 2×2。

在 MobileNetV3 和 XF Block 的基礎上提出了 XFormer，這是一種 CNN-ViT 混合輕量級模型，由堆疊的 MobileNetV3 Block 和 XF Block 組成，用于學習全局和局部上下文信息。與之前設計高效 CNN 的工作一樣，本文的網絡由  塊組成，用于提取原始圖像特征，在特征由 CNN 和transformer blocks處理后，使用 、全局池化和全連接層來產生最終的 logit 預測。

在主要處理塊中，XFormer 有5個階段。前2個階段僅包含 MobileNetV3 Block (MV3) ，因為卷積塊比全 ViT 模型更能提取重要的圖像級特征表示，并隨后幫助轉換器塊看得更清楚。最后三個階段中的每一個都包括一個 MV3 塊和幾個 XF 塊。結合來自 MV3 塊的局部歸納偏差和來自 Transformer Block 的全局信息，網絡可以學習更全面的特征表示，可以輕松地轉移到不同的下游任務。

先前關于 ViT 的工作表明，應在更深的Transformer layers中使用更大的 MLP 比率，并且 Q-K-V 維度應相對小于嵌入維度，以便在性能和模型大小之間進行更好的權衡。作者遵循這些建議并相應地設計輕量級模型。對于三個不同階段的 XF Block ，它們的 MLP 比率、嵌入維度和 Q-K-V 維度分別設置為（2、2、3）、（144、192、240）和（96、96、96）。在全連接層中將 ReLU 替換為 GELU 激活；在所有其他層中使用 SiLU。對于 MV3 Block，擴展比均設置為 4。規格如表 1 所示。

模型的總參數大小只有 550 萬。與類似大小的基于 ViT 的模型相比，本文的模型可以更有效地處理高分辨率圖像并避免潛在的內存瓶頸（見表 2）。

例如，當輸入分辨率為 1024×1024 時，與使用原始自注意力的 MobileViT 相比，XFormer 的推理速度提高了近 2 倍，GPU 內存使用量減少了 32%。本文的模型可以輕松處理高分辨率吞吐量，而不會出現內存瓶頸。最重要的是，XFormer 提供了比比較模型更好的精度（參見表 3），在模型大小和性能之間實現了很好的平衡。

[1].Lightweight Vision Transformer with Cross Feature Attention

本文僅做學術分享，如有侵權，請聯系刪文。