[論文レビュー] CoAtNet: Marrying Convolution and Attention for All Data Sizes

Zihang Dai, Hanxiao Liu+2 more|arXiv (Cornell University)|Jun 9, 2021

Advanced Neural Network Applications参考文献 49被引用数 735

ひとこと要約

CoAtNet は深層分離畳み込みと相対自己注意を統一ブロックで組み合わせ、畳み込みとトランスフォーマーブロックを垂直に積み上げることで、データ域ごとに強い一般化能力と高い容量を実現し、多様なデータ/計算予算の下で最先端の ImageNet 結果を達成します。

ABSTRACT

Transformers have attracted increasing interests in computer vision, but they still fall behind state-of-the-art convolutional networks. In this work, we show that while Transformers tend to have larger model capacity, their generalization can be worse than convolutional networks due to the lack of the right inductive bias. To effectively combine the strengths from both architectures, we present CoAtNets(pronounced "coat" nets), a family of hybrid models built from two key insights: (1) depthwise Convolution and self-Attention can be naturally unified via simple relative attention; (2) vertically stacking convolution layers and attention layers in a principled way is surprisingly effective in improving generalization, capacity and efficiency. Experiments show that our CoAtNets achieve state-of-the-art performance under different resource constraints across various datasets: Without extra data, CoAtNet achieves 86.0% ImageNet top-1 accuracy; When pre-trained with 13M images from ImageNet-21K, our CoAtNet achieves 88.56% top-1 accuracy, matching ViT-huge pre-trained with 300M images from JFT-300M while using 23x less data; Notably, when we further scale up CoAtNet with JFT-3B, it achieves 90.88% top-1 accuracy on ImageNet, establishing a new state-of-the-art result.

研究の動機と目的

• ConvNet と Transformer を組み合わせることで、それぞれの帰納的バイアスと高容量の両方を活用できる理由を動機づける。
• 深さ方向畳み込みと相対自己注意を統合する統一ブロックを提案する。
• 一般化と容量のバランスをとる垂直方向のネットワーク配置を探る。
• データ域（1K、21K、そして大規模 JFT）全体で最先端の性能を実証する。
• 設計選択とレイアウト決定を正当化するアブレーションを提供する。

提案手法

• 深さ方向畳み込みを相対自己注意と統合した単一の Transformer 風ブロックを採用する（事前正規化された相対注意とグローバル/ノイズ付き静的カーネル）。
• すべての位置対 i-j に対してスカラーの相対カーネル w を用いて、追加パラメータなしに効率的なグローバル受容野を実現する。
• 異なる垂直レイアウト（畳み込み重視の段 vs トランスフォーマー重視の段）を比較して、一般化と容量を研究する。
• 初期に局所パターンを捉え、後でグローバルな文脈を得るように設計されたマルチステージ配置（S0 ステム、S1 MBConv、S2 MBConv、S3 Transformer Rel、S4 Transformer Rel）を使用する。
• 大規模事前学習（ImageNet-21K、JFT-300M/3B）とファインチューニングを評価し、ViT および ConvNets に対するデータと計算効率の優位性を示す。

実験結果

リサーチクエスチョン

• RQ1 畳み込みと自己注意を過度なコストなしに単一の計算単位に統合するにはどうすればよいか？
• RQ2 限られたデータでの一般化を最大化し、豊富なデータでの容量を最大化する畳み込みと注意ブロックの垂直配置は何か？
• RQ3 視覚タスクにおいて、相対注意は標準的な注意と比較して一般化と転移性を改善するか？
• RQ4 ハイブリッドなConvNet-TransformerモデルはImageNet-1K、ImageNet-21K、JFT のデータ域全体で最先端の性能を達成できるか？

主な発見

モデル	評価サイズ	パラメータ数	FLOPs	ImageNet Top-1 精度
CoAtNet-2	224x224	75M	15.7B	84.1%
CoAtNet-3	224x224	168M	34.7B	84.5%
CoAtNet-2	384x384	75M	49.8B	85.7%
CoAtNet-3	384x384	168M	107.4B	85.8%
CoAtNet-4	384x384	275M	189.5B	-
CoAtNet-7	512x512	2.44B	2586B	90.88%

• 事前正規化された相対注意ブロックは深さ方向畳み込みと自己注意を効果的に統合し、両方の望ましい特性を保持する単一の計算単位を実現する。
• 畳み込み段を Transformer 段の前に置く垂直積み重ね（例: C-C-T-T）は、よりトランスフォーマー重視の配置と比較して、一般化が向上し、容量も競争力がある。
• 相対注意は標準注意より一般化を改善し、特に転移シナリオ（ImageNet-21K から ImageNet-1K）で顕著である。
• CoAtNet は 1K データのみで ImageNet-1K の top-1 が 86.0%、ImageNet-21K で事前学習しファインチューニングした場合は 88.56% の top-1 を達成し、より大きなデータサイズで学習した ViT-Huge に匹敵する。
• 大規模事前学習（JFT-3B）では ImageNet で 90.88% の top-1 精度を達成し、同程度の計算で新しい最先端を確立。
• アブレーションは、より多くの畳み込み段が一般化を改善し、最適なレイアウトは転送性と効率のために S2 MBConv と S3 Transformer ブロックのバランスを取ることを示す。