なぜビジョントランスフォーマーはこれほど高性能なのか。

Transformer 2022年05月16日

3つの要点
✔️ ViTは、すべての層でより均一な表現（特徴量）を持っている。つまり各層での表現が似ている。
✔️ ViTは、自己注意（self-attention）により早期にグローバルな情報を集約できる。
✔️ ViTは、表現を下位層から上位層へ強く伝搬させる。

Do Vision Transformers See Like Convolutional Neural Networks?
written by Maithra Raghu, Thomas Unterthiner, Simon Kornblith, Chiyuan Zhang, Alexey Dosovitskiy
(Submitted on 19 Aug 2021 (v1), last revised 3 Mar 2022 (this version, v2))
Comments: Published on arxiv.
Subjects: Computer Vision and Pattern Recognition (cs.CV); Artificial Intelligence (cs.AI); Machine Learning (cs.LG); Machine Learning (stat.ML)

code：

本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。

Abstruct

畳み込みニューラルネットワーク（convolutional neural network, CNN）は、これまで視覚データに対するデファクトモデルでした。しかし近年ではViT（Vision Transformer）が画像分類タスク分野を席巻しています。

なぜViTはそれほど画像タスクで良い結果を残せたのでしょうか。本論文ではViTとCNNの内部表現を解析し、2つのアーキテクチャにどのような違いがあるのかを検索しました。

その結果、ViTのセルフアテンション（self-attention, 自己注意）機構によりグローバルな表現を早期に集約できる点と、下位層から上位層へ表現を強く伝搬する残差結合（スキップコネクション）が重要な役割を果たすことが分かりました。

Introduction

画像タスクにおいて過去数年間はCNNが優勢でした。これは畳み込みによる空間的な特徴の保存という再帰的な手法によるところが大きいです。なおかつ容易に転移学習に利用可能で、汎用的な視覚表現を使い回すことができました。しかし最近の研究ではViTが画像分類タスクでCNNを凌駕することが報告されています。ViTは、自然言語処理に用いられるTransformerの自己注意（self-attention）機構を使うことでグローバルな特徴を集約することが可能です。これはCNNが畳み込みという帰納的なアプローチで情報を集約することと対照的で、これまで報告されてきたCNNの改善方法とは異なります。

本論文ではViTがどのように画像タスクをこなしているのかについて以下のように解析します。

ViTは、ResNetよりも下位層でグローバルな特徴を取り込むため、全体として異なる特徴量を有する。
しかしながら下位層での局所的な特徴量の取り込みはやはり重要であり、大規模データで学習した際にViTでも下位層でのアテンションがこれらを学習している。
ViTのスキップ結合はResNetよりも影響力が強く、モデルの性能と表現力に多大な影響を与えている。
分類タスクではなく、物体検知にViTを使用することを考え、入力された空間情報がどの程度保存されるのかを解析した。
データセットの規模が転移学習に及ぼす影響について、線形プローブ法を用いて検討し、高品質な中間表現に対する重要性を明らかにした。

Related Work

画像タスクに対するTransformerの開発は活発な研究分野です。先行研究ではCNNとアテンションを組合せたり、画像サイズを小さくすることでアテンションがどのように局所の特徴を捉えているのかを解析しています。ところがViTのアーキテクチャに工夫を加えるという研究は少なく、CNNとの比較検討はさらに少なくなっています（文章タスクでは報告されてきています）。

背景と実験設定

本研究ではCNNとViTを比較することで、画像タスクを解決する方法に違いがあるのかを検討します。ここではCNNとしてResNetを使用します。特に指定しない限り、データセットはJFT-300Mを用いています。詳細は付記Aを参照のこと。

表現類似性とCKA

隠れ層の表現解析には中心カーネルアライメント（centered kernel alignment）を用います。CKAは、同一画像に対する2つの層の表現（活性化行列）を入力とします（XとY）。ここで以下のようにKとLを定義します。

これをグラム行列といいます。CKAは以下のように計算されます。

HSICとはHilbert-Schmidtの独立性検定と呼ばれる指標です。

図1。CKAをプロットしたものが上図です。ViTは、ViT-L/16とViT-H/14の2種類を使用しています。ResNetも同様です。まず対角線を見てみましょう。同じ層を比較すればCKAが最大になるので、ViTとResNet関係なく対角線は白色です。

そしてViTでは全体的にオレンジ色〜白色になっていることが分かります。これはつまり下位層と上位層が出力する表現が似ていることを意味します。一方、ResNetの方を見るとプロット図がオレンジ色と紫色のチェック柄のようになっています。これはある層と別の層の表現が似ていなかったり（紫色）、似ていたり（オレンジ色）と層ごとの類似性がまだらになっているためです。

このことからViTが獲得した表現性は下位層と上位層で似ていることが明らかになります。そもそもニューラルネットワークは上位層に向かうにつれ、ローカル（局所）からグローバル（大域）な特徴を学習します。そのためViTは初期から大域的な特徴を捉えていると考えることができ、逆に言えば、ViTは層を重ねても異なる表現を獲得しないとも言えます。

図2。図1とは異なり、ViTとResNetを軸として表現の類似性をプロットしています。ViTとResNetでは、ViTの0-20層とResNetの0-60層で類似性が高いように見えます。ViTの40-60層とResNetの80-120層でも同様です。逆に2つのモデルの上位層には類似性が見られません。

図1の内容と合わせて考えると、ViTとResNetの画像の抽象化が異なる様式であることが示唆されます。

層表現におけるローカルとグローバルの情報

self-attentionレイヤーはCNNとは構造的に大きく異なります。self-attentionレイヤーは、マルチヘッドアテンションから構成されていて、それぞれのヘッドについてクエリとアテンションの距離を計算することができます。これによってセルフアテンションレイヤーがどの程度局所的な情報と全体的な情報を集約しているかが分かります。

上記が、この距離の5000データにおける平均値をプロットしたものです。項目は2種類あり、下位層（block0, block1）と上位層（block22, block23）があります。これを見ると下位層はMean Distanceが小さくてかつ幅がある、つまり局所と大域のどちらの情報も取得しており、一方で上位層はほとんど大域的な情報を取得していることが分かります。

これらの結果を総合すると、(i) ViTの下位層はResNetの下位層とは異なる表現力をもち、(ii) ViTは下位層と上位層の間で強く表現が伝搬されており、(iii) ViTの最高層はResNetとは全く異なる表現を持っていることが示唆されました。

また上図に示すように、ViTにおいて学習データセットを小さくする（ImageNetを使う）と、下位層でも局所的な学習が行われていないことが明らかになりました。このときモデルのパフォーマンスは低下し、CNNよりも性能が落ちています。つまりこの場合にはCNNが局所的な特徴を学習することがパフォーマンスにつながっていると考えられます。

さらにResNetと比較したのが上図です。ResNetの下位層とViTの下位層が獲得した表現の類似度は、横軸である平均距離が大きくなるほど、単調に減少することから、距離が大きくなるViT上位層とResNet下位層では定量的に異なる特徴を学習していることが示されています。

スキップ結合を介した表現の伝搬

これまでにViTの表現が均一であることが示されてきましたが、どのようにして表現が伝搬しているのでしょうか。

ここでスキップ結合をしないViTを実験的に訓練しました。すると上図に示すように、表現の類似度がなくなっている（紫〜黒色）ことが分かります（加えてモデルの性能も低下しました）。ViTにおけるスキップコネクションが重要であることの証明です。

空間情報と局在化

上図は、ViTとResNetがどの程度空間情報を保持しているかを示したものです。任意の層の特徴マップでCKA類似度を算出しています。ViT（上段2つ）では概ね似た位置に黄色あるいは青色がきています。ResNetも概ね似た位置の情報を保持しているように見えますが、空間的に幅広く類似性を示しています。つまり関係のない位置でも類似性があることが分かります。

このことからViTの上位層でも空間的な情報は保持されていると考えることができ、物体検知タスクにおいても有用であることが示唆されます。

スケールの影響

これまでにデータセットの大きさによって、ViTが学習する特徴が変化することを示しました。

上記は、データセットの大きさを変化させたときのモデルの性能変化を示しています。左は、JFT-300Mで学習したモデル（実線）とImageNetで学習させたモデル（破線）の比較です。タスクはImageNetでの分類タスクです。

最終的な（横軸で 1.0）結果には大きな差はありませんが、中間層表現（横軸で0.5あたり）を見るとViTの精度が高いことが分かります。これはつまり大規模データ（JFT-300M）で学習した方が、ViTはより高品質な中間表現を獲得できることが示唆されています。

右ではResNetを加えての比較です。最終的な結果自体は、ViTとResNetに大きな差はありませんが、やはり中間表現においてはViTの方が優勢です。またこのことはCIFAR-10とCIFAR-100でも同様の結果でした。

Discussion

CVのブレイクスルーであるCNNに比較して、自然言語から発生したTransformerが画像タスクにおいてもCNNと同様のパフォーマンスを発揮していることは驚嘆に値します。

本論文ではCNN（ResNet）とViTを比較して、内部構造に驚くほど明確な違いがあることを示しました。ViTは初期からグローバルな情報を集約し、それを上位層へスキップコネクションを使って強く伝搬していることが分かりました。さらにViTには大きな事前学習データセットが必要であり、これによって高品質な中間表現を獲得していることが分かりました。これらの知見は、ViTとCNNの違いに関する疑問に答えており、今後の研究の新しい方向性を示すものです。