89.77%！谷歌大腦Quoc V.Le團隊提出CoAtNet：將卷積與自注意力納入同一模塊。

點擊上方“AI算法與圖像處理”，加"星標"或“置頂”
重磅干貨，第一時間送達

paper：https://arxiv.org/abs/2106.04803

本文是谷歌研究院Quov V.Le團隊在卷積與自注意力組合方面的探索，將深度卷積與自注意力集成統(tǒng)一到一個計算模塊中，并從實驗角度對卷積與自注意力的組合方式進行了論證，進而確定了CoAtNet的架構。所提方案在ImageNet數(shù)據(jù)集上取得了大幅超越其他ConvNet與Transformer的性能。比如，無需額外數(shù)據(jù)，CoAtNet在ImageNet上取得了86%的top1精度；額外引入JFT預訓練后，模型進一步提升提升到89.77%，超越了之前最佳的EfficientNetV2與NFNet。

Abstract

Transformer在計算機視覺領域受到 了越多越多的關注，但他們的性能仍落后于優(yōu)秀的CNN。在這篇文章中，我們將表明：盡管Transformer具有非常大的模型容量，但由于缺乏正確的歸納偏置導致其泛化性能不如CNN。

為有效取兩者之長，我們提出了CoAtNet，它基于以下兩個關鍵點而構建的混合模型：

• 深度卷積與自注意力可以通過簡單的相對注意力進行統(tǒng)一；
• 垂直疊加卷積與注意力層對于提升泛化性能、容量以及效率非常有效。

實驗結果表明：在不同數(shù)據(jù)集、不同資源約束下，所提CoAtNet均取得了SOTA性能。比如，無需額外數(shù)據(jù)，CoAtNet在ImageNet上取得了86.0%top1精度；當引入額外的JFT后其性能可以進一步提升到89.77%。更值得注意的是，當在ImageNet21K上預訓練后，CoAtNet可以達到88.56%top1精度，這與JFT上預訓練的ViT-huge精度相當，而CoAtNet需要的訓練數(shù)據(jù)比ViT-huge少23倍。

Method

接下來，我們將聚焦于“如何最優(yōu)化組合卷積與Transformer？”粗略的講，我們將該問題分為兩部分：

• 如何將卷積與自注意力組合到一個基礎計算模塊中？
• 如何垂直堆疊不同類型的計算模塊以構建一個完整網(wǎng)絡？

Merging Convolution and Self-Attention

對于卷積而言，我們主要聚焦于MBConv，它采用深度卷積捕獲空間交互關系。進行該選擇的一個關鍵原因：Transformer中的FFN與MBConv均采用了“Inverted Bottleneck”的設計思想(先對輸入通道數(shù)擴張，然后在收縮到原始通道數(shù)以支持殘差連接)。

除了“Inverted Bottleneck”的相似外，我們還注意到：深度卷積與自注意力均可以表達成預定義感受野范圍內(nèi)值的加權和。具體來說，卷積依賴于固定核在局部感受野內(nèi)收集信息：

作為對比，自注意力使得感受野覆蓋到整個空間位置并基于成對數(shù)據(jù)計算權值后加權：

在正式介紹如何對其進行最佳組合之前，我們先來看一下兩者之間的相對優(yōu)弱所在，這有助于我們理解需要保留的優(yōu)異屬性。

• 首先，對于深度卷積來說，它的卷積核是輸入無關的靜態(tài)值；而對于自注意力而言，它的注意力權值是輸入相關的動態(tài)值。因此，自注意力更易于捕獲不同位置間復雜的相關性，而這個特性是我們在處理高級概念時所需要的。然后，自注意力的這種特性也帶來的過擬合問題，尤其當數(shù)據(jù)非常有限時；
• 其次，注意到：對于任意位置對,卷積的權值僅關注了相關偏移，而未關注特定的值。這種屬性一般特指平移不變性，它有助于提升小數(shù)據(jù)集的泛化性。由于引入了絕對位置嵌入，Transformer并不具備這種特性。這就是為什么小數(shù)據(jù)集上ConvNet的性能通常優(yōu)于Transformer的原因。
• 最后，感受野的尺寸是卷積與自注意力的最關鍵區(qū)別。一般來講，更大的感受野可以提供更多的上下文信息，進而導致更高的模型容量。因此，全局感受野是在視覺領域使用自注意的一個關鍵驅(qū)動力。然而，更大的感受野同時也帶來了更多的計算量。

基于上述比較，一個理想的模型應當能組合上表中的三個屬性。類似深度卷積與自注意力，一個最直接的方式：全局靜態(tài)卷積核與自適應注意力矩陣組合，即：

有意思的是，盡管這個想法看起來過于簡單，但是預歸一化版本對應了相對自注意力的一個特定變種。這種情況下，注意力權值由平移不變權值與輸入自適應聯(lián)合確定。更重要的是，為引入全局卷積核且不會導致大量的參數(shù)量，我們重新加載為標量而非向量，進而只會引入非常少的計算量。接下來，我們將采用帶預歸一化相對注意力變種的Transformer作為CoAtNet的核心模塊。

Vertical Layout Design

前面找到了一種組合卷積與注意力的簡單方法，接下來，我們將考慮如何進行堆疊以構建一個完整網(wǎng)絡。

正如前面所提到：全局上下文會帶來大量的計算量，與空間尺寸成二次關系。如果我們直接采用上述相對注意力到原始輸入圖像，計算效率會非常低。因此，為構建一個實際可行的網(wǎng)絡，我們主要以下三種候選方案：

• A: 執(zhí)行下采樣以降低空間尺寸，在達到可接受水平后采用全局相對注意力；
• B: 采用局部注意力以全局，類似卷積約束自注意力的感受野；
• C: 采用特定的線性注意力替換二次Softmax注意力，進而將計算復雜度降低到與空間尺寸成線性關系。

我們對方案C進行了簡單的實驗，但并未得到一個好的結果。對于方案B，局部注意力的實現(xiàn)會涉及大量的形變變換操作，進而需要大量的內(nèi)存訪問。在TPU上，這種操作的計算效率非常低，這不僅與全局注意力的加速相悖，同樣會影響模型容量。因此，接下來，我們主要聚焦于方案A。

對于方案，下采樣可以通過以下方式得到：

• 類似ViT，采用卷積下采樣到stride=16；
• 類似ConvNet，采用多階段逐漸池化的網(wǎng)絡.

基于上述選擇，我們設計了一個包含5個變種的搜索空間并通過可控實驗進行了對比。

• 當采用ViT Steam時，我們直接堆疊L個Transformer，表示為
• 當采用多階段ConvNet，我們模擬ConvNet構建一個包含5個階段(S0,S1,S2,S3&S4)的網(wǎng)絡，空間分辨率逐漸從S0下降到S4。S0為簡單的兩層卷積，S1采用MBConv與SE。從S2到S4，我們同時考慮MBConv與Transformer ，此時有這樣一個約束：卷積要位于Transformer之前。這就引出了4個變種：C-C-C-C,C-C-C-T, C-C-T-T, C-T-T-T。

為系統(tǒng)研究這些設計選擇，我們主要考慮了泛化性與模型容量兩個引入：

• 泛化性：我們主要感興趣訓練損失與驗證精度之間的差異。如果兩個模型具有相同的訓練損失，具有更高驗證精度的模型具有更好的泛化性。當訓練數(shù)據(jù)比較小時，泛化性對于提高數(shù)據(jù)效率非常重要。
• 模型容量：我們用來評價模型對于大數(shù)據(jù)的擬合能力。當訓練數(shù)據(jù)非常大時(幾乎不可能存在過擬合)，具有更高精度的模型通常具有更大的容量。

為比較模型泛化性與容量，我們在ImageNet與JFT上分別訓練300與3epoch不同變種的混合模型，均未添加任何正則與增廣，訓練損失與驗證精度對比見下圖。

• 對于ImageNet上的結果，一個重要發(fā)現(xiàn)：就泛化性能而言，有以下結論：

也就是說，的泛化性最差。我們認為：這與ViT的一次性下采樣缺乏底層信息有關。

• 在模型容量方面，其排名如下：

這意味著：具有更多的Transformer模塊并不意味著更高的容量。為在前兩者之間做一個決策，我們考慮另外一個“遷移能力”：在JFT上預訓練，然后在ImageNet微調(diào)。結果見下表，我們最終選擇了配置。

考慮到泛化性能、模型容量、遷移性能以及計算效率，我們采用了構建CoAtNet，其架構見上圖。

Experiments

接下來，我們在相對公平的設置下對所提CoAtNet與其他方案進行了對比。下表給出了不同配置CoAtNet的參數(shù)信息。

Main Results

ImageNet-1K 上表給出了僅在ImageNet上訓練的模型性能對比，從中可以看到：在相似前提下，CoAtNet不僅優(yōu)于ViT變種，同時與最佳卷積方案(比如EfficientNetV2、NFNet)的性能相媲美。下圖Figure2給出了僅在ImageNet上訓練時模型結果的可視化圖，可以看到CoAtNet具有比其他注意力模型更佳的性能。

ImageNet-21K 如上述Table4與Fig3所示，當采用ImageNet21K預訓練后，CoAtNet的優(yōu)勢更明顯，大幅優(yōu)于其他所有模型。比如，CoAtNet取得了88.56%的精度，與ViT-H/14的88.55%相媲美，同時所需額外的預訓練數(shù)據(jù)更少。

JFT 上表對比了采用JFT預訓練時的模型性能對比，可以看到：CoAtNet-4取得了與之前NFNet-F4相當?shù)男阅?、相當?shù)膮?shù)量，同時具有更快的TPU訓練速度；而CoAtNet-5則進一步達到了89.77%的top1精度，優(yōu)于其他同等配置的模型。

全文到此結束，更多消融實驗與分析建議查看原文。

個人微信（如果沒有備注不拉群！）
請注明：地區(qū)+學校/企業(yè)+研究方向+昵稱



下載1：何愷明頂會分享

在「AI算法與圖像處理」公眾號后臺回復：何愷明，即可下載。總共有6份PDF，涉及 ResNet、Mask RCNN等經(jīng)典工作的總結分析

下載2：終身受益的編程指南：Google編程風格指南

在「AI算法與圖像處理」公眾號后臺回復：c++，即可下載。歷經(jīng)十年考驗，最權威的編程規(guī)范！

下載3 CVPR2021

在「AI算法與圖像處理」公眾號后臺回復：CVPR，即可下載1467篇CVPR 2020論文 和 CVPR 2021 最新論文