paper: https:///abs/2107.12292

code: https://github.com/JDAI-CV/CoTNet

本文是京東AI研究院梅濤團隊在自注意力機制方面的探索，不同于現(xiàn)有注意力機制僅采用局部或者全局方式進行上下文信息獲取，他們創(chuàng)造性的將Transformer中的自注意力機制的動態(tài)上下文信息聚合與卷積的靜態(tài)上下文信息聚合進行了集成，提出了一種新穎的Transformer風格的“即插即用”CoT模塊，它可以直接替換現(xiàn)有ResNet架構(gòu)Bottleneck中的卷積并取得顯著的性能提升。無論是ImageNet分類，還是COCO檢測與分割，所提CoTNet架構(gòu)均取得了顯著性能提升且參數(shù)量與FLOPs保持同水平。比如，相比EfficientNet-B6的84.3%，所提SE-CoTNetD-152取得了84.6%同時具有快2.75倍的推理速度。

Abstract

Transformer極大的促進了NLP領域的發(fā)展，Transformer風格的網(wǎng)絡架構(gòu)近期在CV領域也掀起了一波浪潮并取得了極具競爭力的結(jié)果。盡管如此，現(xiàn)有的架構(gòu)直接在2D特征圖上執(zhí)行自注意力并得到注意力矩陣，但是關(guān)于近鄰豐富的上下文信息卻并未得到充分探索。

本文設計了一種新穎的Transformer風格模塊CoT(Contextual Transformer, CoT)用于視覺識別，該設計充分利用輸入的上下文信息并引導動態(tài)注意力矩陣的學習，進而提升了視覺表達能力。技術(shù)上來講，CoT模塊首先通過卷積對輸入keys進行上下文信息編碼得到關(guān)于輸入的靜態(tài)上下文表達；進一步將編碼keys與輸入query進行拼接并通過兩個連續(xù)卷積學習動態(tài)多頭注意力矩陣；所得注意力矩陣與輸入values相乘即可得到關(guān)于輸入的動態(tài)上下文表達。靜態(tài)上下文表達與動態(tài)上下文表達的融合結(jié)果作為該模塊的輸出。所提CoT模塊可以作為“即插即用”模塊替換ResNet中的卷積并得到Transformer風格的架構(gòu)，故而稱之為CoTNet(Contextual Transformer Network)。

多個領域(包含圖像分類、目標檢測、實例分割)的充分實驗結(jié)果表明：CoTNet是一種更強的骨干網(wǎng)絡。比如，在ImageNet分類任務中，相比ResNeSt101，CoTNet取得了0.9%的性能提升；在COCO目標檢測與實例分割任務中，相比ResNeSt，CoT分別取得了1.5%與0.7%mAP指標提升。

Method

我們首先對現(xiàn)有骨干網(wǎng)絡中采用的傳統(tǒng)自注意力模塊進行簡單介紹；然后引入了本文所設計的Transformer風格的模塊；最后再引入基于所提CoT模塊構(gòu)建的CoTNet、CoTNeXt等等架構(gòu)。

Multi-head Self-attention in Vision Backbones

上面Figure-a給出了現(xiàn)有骨干中的多頭自注意力模塊架構(gòu)圖。具體來說，給定2D輸入特征X(尺寸為)，我們通過如下公式將其變換為queries、keys以及values：

注：這里采用的是卷積。然后我們在K與Q之間計算局部相關(guān)矩陣R：

局部相關(guān)矩陣R進一步通過位置信息進行豐富化：

接下來，通過對上述局部相關(guān)矩陣進行歸一化即得到注意力矩陣A :

最后，通過values與注意力矩陣A相乘即得到增強后特征Y：

Contextual Transformer Block

傳統(tǒng)的自注意力僅在空域進行信息交互，依賴于輸入自身，即相關(guān)性通過獨立的方式學習所得，而忽視了近鄰間豐富的上下文信息。這無疑會嚴重限制2D特征圖的自注意力學習能力。為緩解該問題，我們構(gòu)建了一種新的Transformer風格的模塊CoT(見上面的Figure2-b)，它將上下文信息挖掘與自注意力學習集成到統(tǒng)一架構(gòu)中。本文的出發(fā)點：充分探索近鄰間的上下文信息以一種高效方式提升自注意力學習，進而提升輸出特征的表達能力。

具體來說，假設輸入2D特征，keys, queries以及values分別定義為。不同于傳統(tǒng)自注意力采用卷積編碼keys，CoT模塊首先采用組卷積提取上下文信息，所得自然的可以反應近鄰間的上下文信息，我們將其視作輸入X的靜態(tài)上下文表達。

然后，將前述所得keys與Q進行拼接并通過兩個連續(xù)卷積計算注意力矩陣：

換句話說，每個位置局部相關(guān)矩陣基于query與keys學習所得，而非獨立的query-key對。這種處理方式可以通過額外的靜態(tài)上下文提升自注意力學習能力。接下來，基于上述所得注意力矩陣A，我們生成增強特征:

上述所得增強特征可以捕獲關(guān)于輸入的動態(tài)特征交互，我們將其稱之為輸入X的動態(tài)上下文表達。最后，我們將上述兩種上下文表達通過注意力機制融合(比如SKNet中的注意力融合機制)即可得到CoT模塊的輸出。

Contextual Transformer Networks

所提CoT的設計是一種統(tǒng)一的注意力模塊，它可以作為ConvNet中標準卷積的替代。因此，我們可以將其替換現(xiàn)有ConvNet中的卷積進而提升其表達能力。

接下來，我們簡要介紹一下如何將CoT模塊嵌入到現(xiàn)有ResNet架構(gòu)中且不會顯著提升參數(shù)量。下表給出了將CoT模塊嵌入到ResNet與ResNeXt中的信息示意，并將所得模型稱之為CoTNet-50、CoTNeXt-50。當然，還可以將CoT模塊嵌入到ResNet101中，更詳細的嵌入方式可參考官方code。

具體來說，CoTNet-50直接采用CoT替換Bottlenck中的卷積；類似的，CoTNeXt-50采用CoT模塊替換對應的組卷積，為獲得相似計算量，對通道數(shù)、分組數(shù)進行了調(diào)整：CoTNeXt-50的參數(shù)量是ResNeXt-50的1.2倍，F(xiàn)LOPs則是1.01倍。

Experiments

接下來，我們通過不同領域(包含ImageNet分類、COC目標檢測以及實例分割)的實驗評估所提CoTNet的性能。

ImageNet Classification

上表給出了ImageNet默認配置訓練的性能對比，可以看到：

• 相比已有ResNet改進，在同等參數(shù)量與FLOPs約束下，所提CoTNet取得了最佳性能；
• 采用局部自注意力的LRNet-50與Stand-Alone取得了比ResNet50更佳的性能；
• 采用了全局自注意力的AA-ResNet-50與LambdaResNet-50取得了進一步的性能提升，但仍低于SE-ResNet；
• 同時利用局部與全局上下文信息的CoTNet取得了優(yōu)于SE-ResNet的性能，這說明：上下文信息挖掘與自注意力學習的集成是一種有效的增強表達能力的方式。

上表給出了更先進訓練配置下的性能對比，從中可以看到：

• 所提CoTNet、CoTNeXt均具有比其他ResNet改進版更優(yōu)的性能；
• 相比ResNeSt-50，ResNeSt-101，所提CoTNeXt-50與CoTNeXt-101分別取得了1.0%與0.9%的性能提升；
• 相比BoTNet，所提CoTNet同樣具有更優(yōu)的性能；甚至于，SE-CoTNetD-152(320)取得了比BoTNet-S1-128(320)、EfficientNet-B7更優(yōu)的性能.

Object Detection

上表給出了COCO目標檢測任務上的性能對比，從中可以看到：相比ResNet、ResNeSt等骨干，所提CoTNet取得了顯著的性能提升。這說明：上下文信息挖掘與自注意力學習的集成是一種有效的特征增強方法，甚至可以很好的遷移到下游任務中。

Instance Segmentation

上表給出了COCO實例分割任務上的性能對比，從中可以看到：相比ResNet、ResNeSt等骨干，所提CoTNet取得了顯著的性能提升。

Ablation Study

上面兩個圖給出了精度-推理耗時的對比，從中可以看到：相比其他骨干網(wǎng)絡，CoTNet取得了最佳的top1精度，同時具有更少的推理耗時。也就是說，CoTNet具有更好的精度-速度均衡。比如，相比EfficientNet-B6，所提SE-CoTNetD-152(320)不僅精度高0.6%，推理速度還快2.75倍。

上表對比了不同上下文信息獲取的性能對比，從中可以看到：

• 靜態(tài)上下文方式(可視作無自注意力的ConvNet)僅取得77.1%的top1精度；
• 動態(tài)上下文方式可以取得78.5%的top1精度；
• 動靜上下文簡單的相加融合可以取得78.7%的top1精度；
• 通過注意力融合的方式集成動靜上下文可以進一步將其性能提升到79.2%。

上表對比了不同階段替換CoT以及SE注意力等改進的性能影響，從中可以看到：

• 提升替換CoT模塊的數(shù)量可以有效提升模型的性能，而參數(shù)量與FLOPs不會顯著變化；

• 相比SE-ResNetD-50，所提SE-CoTNetD-50取得了更佳的性能。

如果覺得有用，就請分享到朋友圈吧！