Do You Even Need Attention? A Stack of Feed-Forward Layers Does Surprisingly Well on ImageNet

[e4exp]

• https://arxiv.org/abs/2105.02723
• https://github.com/lukemelas/do-you-even-need-attention
• 2021

視覚変換器の注目層をフィードフォワード層に置き換えてみたところ、ImageNetでもうまく機能することがわかりました。

画像分類などの視覚タスクにおいてビジョントランスフォーマーが優れた性能を発揮するのは、マルチヘッドアテンションレイヤーの設計によるものであることが多い。 しかし、注意がどの程度、この強力なパフォーマンスに関与しているのかは、まだ明らかになっていません。 このレポートでは、「注意層は必要なのか」を問いかけます。 具体的には、視覚変換装置の注意層を、パッチ次元に適用されるフィードフォワード層に置き換えます。 結果として、パッチ次元と特徴次元に交互に適用された一連のフィードフォワード層だけのアーキテクチャとなります。 ImageNetを用いた実験では、このアーキテクチャは驚くべき性能を発揮しました。 ViT/DeiTベースのモデルは74.9%のトップ1精度を獲得しましたが、ViTでは77.9%、DeiTでは79.9%でした。 これらの結果は、パッチ埋め込みなど、注意以外の視覚変換の側面が、従来考えられていたよりも高い性能を発揮している可能性を示しています。 これらの結果を受けて、現在のモデルがなぜ効果的なのかを理解するために、コミュニティがより多くの時間を費やすことを期待しています。

これは、Google ResearchのMLP-Mixerと同時進行の研究です。アイデアは全く同じですが、1つの違いは、彼らはより多くの計算機を使用することです。 https://arxiv.org/abs/2105.01601

[e4exp]

1. はじめに

[4]で紹介されたビジョン変換器のアーキテクチャは，一連の画像パッチに一連の変換器ブロックを適用します。 各ブロックは，マルチヘッドアテンション層[11]と，それに続くフィードフォワード層（線形層または単層MLP）で構成されており，特徴次元に沿って適用されます． このアーキテクチャの汎用性の高さと，画像分類ベンチマークでの高いパフォーマンスにより，ビジョンコミュニティから大きな関心が寄せられている． しかし、ビジョン変換が有効である理由は、まだはっきりとはわかっていません。 視覚タスクにおいてトランスフォーマーが成功した理由として最も挙げられるのは、モデルにグローバルな受容野を与える注目層の設計である。 この層は、データに依存する線形層と見なすことができ、画像パッチに適用すると、畳み込みに似ています（正確には同等ではありませんが）。 実際、最近の研究では、注目層の効率と有効性を向上させるために、かなりの量の研究が行われている。

このレポートでは、そもそもなぜ視覚変換がうまくいくのか、その理由に少しでも光を当てたいと思い、実験を行いました。 具体的には、視覚変換器から注意を取り除き、その代わりにパッチ次元に適用されるフィードフォワード層に置き換えます。 この変更により、モデルは単純に、パッチ次元と特徴次元に交互に適用される一連のフィードフォワード層となります（図1）。

ImageNetでの実験（表1）では、注意を払わなくても非常に高い性能が得られることが示されました。

特に、ViTベースのモデルは、ハイパーパラメータの調整なしで（つまり、ViTモデルと同じハイパーパラメータを使用して）、74.9%のトップ1精度を達成した。 これらの結果は、ビジョントランスフォーマーの高い性能は、注意メカニズムよりも、パッチ埋め込みや慎重に選択されたトレーニング補強セットによって生じる誘導バイアスなど、他の要因に起因する可能性を示唆している。 このレポートの主な目的は、シンプルなアーキテクチャの限界を探ることです。 ImageNetベンチマークを破ることではありません。

この分野では、ニューラルアーキテクチャ検索（例：EfficientNet[8]）のような手法が必然的に最も良い結果を出すでしょう。 それにもかかわらず、私たちは、コミュニティがこれらの結果を興味深く感じ、これらの結果がより多くの研究者に、現在のモデルがなぜそのように効果的なのかを調査するよう促してくれることを願っています。

[e4exp]

2. 背景

このレポートの背景には、ここ数ヶ月の間に、視覚変換器アーキテクチャのバリエーションに関する研究が爆発的に増えたことがあります。 Deit [9]は蒸留を追加し、DeepViT [14]は注目ヘッドを混合し、CaiT [10]は注目層を2段階に分離し、Token-to-Token ViT [13]はネットワーク全体で隣接するトークンを集約し、CrossViT [1]は2つのスケールでパッチを処理します。 PiT [6]はプーリング層を追加し，LeViT [5]は畳み込みエンベッディングと修正された注目/正規化層を使用し，CvT [12]は注目層で深さ方向の畳み込みを使用し，Swin/Twins [2, 7]はグローバルな注目とローカルな注目を組み合わせています． これらの作品は，画像変換アーキテクチャを改良したもので，それぞれがImageNetで強力な性能を示しています． しかし，ViTやその多くのバリエーションの異なる部分が，これらのモデルの最終的な性能にどのように寄与しているかは明らかではない． このレポートでは、この問題の1つの側面、すなわちViTの成功に注目層がどれほど重要であるかを調査する実験について詳述する。

[e4exp]

3. 方法と実験

3.1. アテンションは必要か？ ViTモデルから注意層を取り除き、パッチ次元の単純なフィードフォワード層に置き換えます。 つまり、パッチ次元を高次元空間に射影し、非線形性を適用し、元の空間に射影するという、特徴次元における標準的なフィードフォワードネットワークと同じ構造を使用します。 図2は、フィードフォワードのみの変換器の1つのブロックのPyTorchコードです。 視覚変換器やその多くのバリエーションと同様に、このフィードフォワードのみのネットワークは、畳み込みネットワークに強く似ていることに注意しなければなりません。 実際、パッチ次元のフィードフォワード層は、完全な受容野と単一のチャンネルを持つ、珍しいタイプの畳み込みとみなすことができます。 特徴次元のフィードフォワード層は1x1の畳み込みと見ることができるので、技術的にはネットワーク全体が一種の畳み込みネットワークを装っていると言ってもよいでしょう。 とはいえ、伝統的に設計された畳み込みネットワーク（ResNet/VGGなど）よりも、構造的には変圧器に似ています。

3.2. 実験セットアップ DeiT [9]のセットアップを用いて，ImageNet [3]上で，ViT/DeiTのtiny，base，largeの各ネットワークに対応する3つのモデルを学習します。 tinyおよびbaseネットワークはパッチサイズ16，largeネットワークは計算上の制約からパッチサイズ32としました． 学習と評価は，解像度224pxで行った． また，全てのモデルでDeiTと全く同じハイパーパラメータを使用しており，ハイパーパラメータのチューニングによって性能が向上する可能性があることが分かりました．

3.3. 結果 表1はImageNetにおける我々の単純なフィードフォワードネットワークの性能を示しています。 最も注目すべき点は、フィードフォワードのみのバージョンのViT/Deit-baseが、多くの古い畳み込みネットワーク（例：VGG16、ResNet-34）に匹敵する驚くべき性能（74.9%のトップ1精度）を達成したことです。 このような比較は、フィードフォワードモデルがより強力な学習拡張を使用しているため、正確には公正ではありませんが、それにもかかわらず、絶対的な意味で非常に強力な結果です。 大規模なモデルでは、注意を払ったモデルと注意を払わなかったモデルの両方で性能が低下しており、トップ1精度はそれぞれ71.2%と71.4%でした。 4]で詳しく述べられているように、このような巨大なモデルには、より大きなデータセットでの事前トレーニングが必要であると思われる。

3.4. フィードフォワード層は必要か？ フィードフォワード層だけでモデルを学習したのだから、当然、アテンション層だけでモデルを学習してみました。 このモデルでは、特徴次元のフィードフォワード層を、特徴次元のアテンション層に置き換えただけです。 実験したのは4.0Mパラメータの小さなモデルだけでしたが、その性能は驚くほど低く（100エポックでのトップ1精度は28.2％、この時点で実行を終了しました）。

3.5. 考察 以上の実験により、アテンション層なしでもそれなりに強力な変換型画像分類器を学習することが可能であることがわかった。 さらに、フィードフォワード層を持たないアテンション層では、同様に強力な性能は得られないようです。 これらの結果は、ViTの強力な性能は、アテンション層のデザインよりも、パッチエンベッディングと学習手順に起因している可能性を示しています。 特にパッチエンベッディングは強い帰納的バイアスを提供しており、このモデルの高いパフォーマンスの主な要因の1つであると考えられます。 実用的な観点から見ると、フィードフォワードのみのモデルは、ビジョン変換器に比べて1つの顕著な利点があります。 これは、パッチ上に適用されるフィードフォワード層内の中間的な投影次元によるもので、そのサイズは必ずしもシーケンス長には依存しません。 通常、中間次元は入力特徴数（すなわちパッチ数）の倍数になるように選択され、その場合、モデルは確かに二次的になりますが、必ずしもそうである必要はありません。 フィードフォワードのみのモデルの大きな欠点は、性能が低いことの他に、固定長のシーケンスでしか機能しないことです（パッチ上のフィードフォワード層のため）。 これは、画像が標準的なサイズに切り取られる画像分類においては大きな問題ではありませんが、他のタスクへの適用が制限されてしまいます。 フィードフォワードのみのモデルは、視覚変換や注意メカニズム全般に光を当てている。 将来的には、これらの結論が画像領域以外、例えばNLPやオーディオなどにどの程度適用されるかを調査することは興味深いことです。

3.6. 結論 この短いレポートでは、注目層を持たないトランスフォーマースタイルのネットワークが、驚くほど強力な画像分類器になることを示しています。 今後の課題としては、変換器アーキテクチャの他の部分（正規化層や初期化スキームなど）の貢献度をよりよく理解することです。 さらに、この短いレポートが、現在のモデルがなぜ優れた性能を発揮するのかについて、さらなる調査を促すものであることを願っています。

[e4exp]

from torch import nn
class LinearBlock(nn.Module):

    def __init__(self, dim, mlp_ratio=4., drop=0., drop_path=0., act=nn.GELU,
        norm=nn.LayerNorm, n_tokens=197): # 197 = 16**2 + 1
        super().__init__()

        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        # FF over features
        self.mlp1 = Mlp(in_features=dim, hidden_features=int(dim*mlp_ratio), act=act, drop=drop)
        self.norm1 = norm(dim)
        # FF over patches
        self.mlp2 = Mlp(in_features=n_tokens, hidden_features=int(n_tokens*mlp_ratio), act=act, drop=drop)
        self.norm2 = norm(n_tokens)

    def forward(self, x):
        x = x + self.drop_path(self.mlp1(self.norm1(x)))
        x = x.transpose(-2, -1)
        x = x + self.drop_path(self.mlp2(self.norm2(x)))
        x = x.transpose(-2, -1)
        return x

class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features, act_layer=nn.GELU, drop=0.):
        super().__init__()
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, in_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x