論文解説: ViT — 画像を単語のように扱う

ViT原論文（Dosovitskiy et al., 2020/2021）を本文だけを根拠に読む。16×16のパッチを単語として扱うという一手、位置埋め込みのアブレーション、CNNの帰納バイアスを捨てた代償、そして「大規模データでCNNを超える」という条件付きの主張の中身。

An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

画素をそのまま並べると破綻する

自己注意は系列の要素どうしを総当たりで比べます。ならば画像も画素の列にすればよさそうですが、論文は関連研究の冒頭でこれを否定します。素朴に適用すると各画素が他の全画素に注意することになり、画素数に対して二次のコストがかかるためスケールしない。だから当時の研究は局所近傍だけ、疎な近似、軸ごと、といった特殊な注意パターンを設計しており、それらは「アクセラレータ上で効率的に実装するには複雑な工学を要する」（§2）。ここが出発点です。

パッチに切る、という一手

ViTの操作は単純です。画像 $\mathbf{x}\in\mathbb{R}^{H\times W\times C}$ を $P\times P$ の固定サイズのパッチへ切り、平坦化して学習可能な線形射影で $D$ 次元に落とす。パッチ数は $N=HW/P^{2}$ で、これがそのままTransformerの実効系列長になります（§3.1）。

$\mathbf{x}^{i}_{p}\mathbf{E}$ が $i$ 番目のパッチの埋め込み、$\mathbf{x}_{\text{class}}$ はBERTの [class] トークンに倣って先頭に足す学習可能なベクトル、$\mathbf{E}_{pos}$ が位置埋め込み。あとは標準のTransformerエンコーダに通し、最終層の先頭トークンの状態を画像表現とします（§3.1）。

この式は要するに、パッチを一列に並べ、どのパッチにも属さない「画像全体の判定を溜めるための席」を先頭にひとつ足し、それぞれの席に「元の画像のどこから来たか」を記したベクトルを足し込む、ということです。$\mathbf{z}_{0}$ はその一列そのもので、エンコーダはこれを入力として受け取ります。

つまり「16×16の画素ブロック1個＝単語1個」。タイトルはこの読み替えの宣言です。元のTransformerにできるだけ忠実に保つ意図は明示されていて、利点はNLP側のスケーラブルな実装をほぼそのまま流用できること（§3）。ViT-L/16 はLargeをパッチ16で使う意味で、系列長はパッチサイズの2乗に反比例するため、小さいパッチほど計算は高くつきます（§4.1）。

位置埋め込みは、ほとんど何でもよかった

位置埋め込みは学習可能な1次元のものを使います。2次元を意識した高度な埋め込みから有意な向上が観測されなかったから、というのが理由です（§3.1）。付録D.4のアブレーションは端的で、ImageNet 5-shot linear の値は位置埋め込み無し 0.61382、1次元 0.64206、2次元 0.64001、相対位置 0.64032。「有る／無い」の差は大きいが「どう符号化するか」の差はほぼ無い（表8）。入力が画素単位ではなくパッチ単位で空間解像度が $14\times14$ 程度しかないため、どの方式でも同じくらい簡単に学べるのだろう、というのが論文の推測です。

そして §4.5 の可視化では、学習後の位置埋め込みは近いパッチほど似た値を持ち、行と列の構造が現れます。2次元の位相を勝手に学んでいた——これが手作りの2次元埋め込みが効かない理由の説明になる、と論文は書きます。位置の与え方はTransformerの位置エンコーディングも併せてどうぞ。