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

ViT 的最新進(jìn)展在視覺(jué)識(shí)別任務(wù)中取得了出色的表現(xiàn)。卷積神經(jīng)網(wǎng)絡(luò) (CNN) 利用空間歸納偏差來(lái)學(xué)習(xí)視覺(jué)表示，但這些網(wǎng)絡(luò)是空間局部的。ViTs 可以通過(guò)其self-attention機(jī)制學(xué)習(xí)全局表示，但它們通常是heavy-weight的，不適合移動(dòng)設(shè)備。

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

實(shí)驗(yàn)結(jié)果表明，XFormer 在不同的任務(wù)和數(shù)據(jù)集上優(yōu)于眾多基于 CNN 和 ViT 的模型。在 ImageNet-1K 數(shù)據(jù)集上，XFormer 使用 550 萬(wàn)個(gè)參數(shù)實(shí)現(xiàn)了 78.5% 的 top-1 準(zhǔn)確率，在相似數(shù)量的參數(shù)下，比 EfficientNet-B0（基于 CNN）和 DeiT（基于 ViT）的準(zhǔn)確率分別提高了 2.2% 和 6.3%。

XFormer的模型在轉(zhuǎn)移到目標(biāo)檢測(cè)和語(yǔ)義分割任務(wù)時(shí)也表現(xiàn)良好。在 MS COCO 數(shù)據(jù)集上，XFormer 在 YOLOv3 框架中超過(guò) MobileNetV2 10.5 AP（22.7 → 33.2 AP），只有 6.3M 參數(shù)和 3.8G FLOPs。在 Cityscapes 數(shù)據(jù)集上，只有一個(gè)簡(jiǎn)單的 all-MLP ****，XFormer 實(shí)現(xiàn)了 78.5 的 mIoU 和 15.3 的 FPS，超過(guò)了最先進(jìn)的輕量級(jí)分割網(wǎng)絡(luò)。

一個(gè)標(biāo)準(zhǔn)的 ViT 模型首先使用patch size h×w 將輸入  reshape為一系列flattened patches ，其中  和  表示token數(shù)。然后將  投影到固定的 D 維空間  并使用一堆transformer blocks來(lái)學(xué)習(xí)inter-patch表示。由于忽略了空間歸納偏差，ViT 通常需要更多參數(shù)來(lái)學(xué)習(xí)視覺(jué)信息。此外，transformers 中 self-attention 的昂貴計(jì)算導(dǎo)致優(yōu)化此類模型的瓶頸。

在本節(jié)提出了 XFormer，這是高效、輕量級(jí)的 CNN-ViT 框架，以解決 ViT 中的上述問(wèn)題。首先介紹了一種提高自注意力效率的新方法，然后說(shuō)明了新的 CNN-ViT 混合模型的架構(gòu)設(shè)計(jì)。

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

其中 σ 是 softmax 操作， 是head維度。計(jì)算注意力分?jǐn)?shù)的計(jì)算復(fù)雜度為 。 self-attention 的二次復(fù)雜性導(dǎo)致了巨大的計(jì)算瓶頸，這使得 ViT 模型難以在移動(dòng)設(shè)備上按比例縮小。

為了解決 self-attention 中的二次復(fù)雜性問(wèn)題，作者提出了一種新的注意力模塊結(jié)構(gòu)，稱為Cross Feature Attention（XFA）。在之前的工作之后，首先沿特征維度 D 對(duì)query Q 和key K 應(yīng)用 L2 歸一化：

 在原始 self-attention 中的直覺(jué)是定位應(yīng)該關(guān)注的重要圖像塊。

但是直接計(jì)算會(huì)導(dǎo)致不必要的冗余和計(jì)算開(kāi)銷。相反，作者為 K 構(gòu)建了2個(gè)中間分?jǐn)?shù)：查詢上下文分?jǐn)?shù)  和查詢特征分?jǐn)?shù) 。使用2個(gè)卷積核矩陣  和  沿token維度 N 計(jì)算 ，沿特征維度 D 計(jì)算 。借助卷積濾波器中間分?jǐn)?shù)向量可以表示計(jì)算注意力圖的更緊湊的表示，同時(shí)也降低了計(jì)算成本。 和  表示為：

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

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

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

1. XFA 通過(guò)構(gòu)建中間查詢上下文和特征分?jǐn)?shù)，沿特征維度D計(jì)算注意力圖，大大降低了計(jì)算成本；

2. XFA 使用可學(xué)習(xí)的溫度縮放參數(shù)來(lái)調(diào)整歸一化，并且不受 softmax 操作的影響。

本文方法為二次復(fù)雜度問(wèn)題提供了解決方案，并且對(duì)于資源受限的設(shè)備更有效且對(duì)移動(dòng)設(shè)備更友好。

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

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

利用提出的Cross Feature Attention模塊介紹了 XF Block，一個(gè)精心設(shè)計(jì)的輕量級(jí) transformer module。XF Block內(nèi)的操作可以表述為：

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

對(duì)于較大的模型（ViT-Large），基于 ViT 的模型通常采用 8×8、16×16 甚至 32×32 的Patch Size。具有較大Patch Size的優(yōu)點(diǎn)之一在于，對(duì)于分類等圖像級(jí)任務(wù)，ViT 可以有效地提取圖像塊信息，而不會(huì)增加過(guò)多的計(jì)算開(kāi)銷。

最近的工作表明，當(dāng)遷移到語(yǔ)義和實(shí)例分割等下游任務(wù)時(shí)，更小的Patch Size更受青睞，因?yàn)樗梢栽鰪?qiáng) Transformer 學(xué)習(xí)更好的像素級(jí)信息的能力，這通常會(huì)帶來(lái)更好的性能。

此外，隨著Patch Size的減小，token數(shù) N 會(huì)大得多。線性復(fù)雜度 XFA 模塊可以避免潛在的計(jì)算瓶頸。在網(wǎng)絡(luò)設(shè)計(jì)中，每個(gè) XF 塊的Patch Size設(shè)置為 2×2。

在 MobileNetV3 和 XF Block 的基礎(chǔ)上提出了 XFormer，這是一種 CNN-ViT 混合輕量級(jí)模型，由堆疊的 MobileNetV3 Block 和 XF Block 組成，用于學(xué)習(xí)全局和局部上下文信息。與之前設(shè)計(jì)高效 CNN 的工作一樣，本文的網(wǎng)絡(luò)由  塊組成，用于提取原始圖像特征，在特征由 CNN 和transformer blocks處理后，使用 、全局池化和全連接層來(lái)產(chǎn)生最終的 logit 預(yù)測(cè)。

在主要處理塊中，XFormer 有5個(gè)階段。前2個(gè)階段僅包含 MobileNetV3 Block (MV3) ，因?yàn)榫矸e塊比全 ViT 模型更能提取重要的圖像級(jí)特征表示，并隨后幫助轉(zhuǎn)換器塊看得更清楚。最后三個(gè)階段中的每一個(gè)都包括一個(gè) MV3 塊和幾個(gè) XF 塊。結(jié)合來(lái)自 MV3 塊的局部歸納偏差和來(lái)自 Transformer Block 的全局信息，網(wǎng)絡(luò)可以學(xué)習(xí)更全面的特征表示，可以輕松地轉(zhuǎn)移到不同的下游任務(wù)。

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

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

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

[1].Lightweight Vision Transformer with Cross Feature Attention

本文僅做學(xué)術(shù)分享，如有侵權(quán)，請(qǐng)聯(lián)系刪文。