Vision Transformer (ViT)

Vision Transformer（ViT）は、言語モデリング向けのTransformerアーキテクチャを画像に適用するニューラルネットワークです。画像を固定サイズのパッチに分割し、各パッチをトークン埋め込みに変換して、自己注意で系列を処理するため、最初の層から各パッチが他のすべてのパッチの情報を利用できます。Googleの研究者は「An Image is Worth 16x16 Words」でこのアーキテクチャを発表しました。この研究では、十分なデータで事前学習すれば、畳み込み層を含まない純粋なTransformerが、画像分類で最高水準の畳み込みニューラルネットワーク（CNN）に匹敵する性能を達成できることが示されました。現在ViTは、視覚言語モデル、セグメンテーション用の基盤モデル、自己教師あり学習のバックボーンとして使われています。

Vision Transformerの仕組み#

ViTは画像を系列に変換し、テキストと同じように処理します。以下では224×224画像に対するViT-Base/16の処理手順を示します。

1. パッチ埋め込み： 画像を16×16ピクセルのパッチに分割し、14×14のグリッドに並んだ196個のパッチを作成します。各パッチをフラット化して学習可能な行列を乗算し、768次元の埋め込みに変換します。通常、これはストライド16の16×16畳み込みとして実装されます。
2. クラス・トークンと位置埋め込み： 学習可能なクラス・トークンを系列の先頭に追加し、学習可能な位置埋め込みをすべてのトークンに加算します。注意機構だけでは順序を認識できないため、各パッチの位置に関する情報はこれらの埋め込みだけから得られます。
3. Transformerエンコーダー： 同一構造のブロックを積み重ねます。ViT-Baseでは12ブロックを使い、すべてのトークン間で情報を混合するマルチヘッド自己注意と、GELU活性化を使って各トークンを個別に変換する2層MLPを交互に適用します。各サブレイヤーの前にレイヤー正規化を行い、サブレイヤーを残差接続で囲みます。
4. ヘッド： クラス・トークンの最終出力を線形分類器に入力します。セグメンテーションなどの密な予測タスクでは、代わりにパッチ・トークンを使い、14×14のグリッドに戻します。

以下のPyTorchコードでは、1枚の画像からトークン系列を構築し、単一のエンコーダーブロックに入力します。

import torch
from torch import nn

# ViT-Base/16の設定：16x16パッチ、768次元トークン、12個の注意ヘッド
image = torch.randn(1, 3, 224, 224)
patch_embed = nn.Conv2d(3, 768, kernel_size=16, stride=16)  # パッチごとに線形射影を1つ適用

# 画像を14x14のパッチグリッドに分割し、フラット化して196個のトークンに変換
tokens = patch_embed(image).flatten(2).transpose(1, 2)

# 学習可能なクラス・トークンを先頭に追加し、学習可能な位置埋め込みを加算
cls_token = nn.Parameter(torch.zeros(1, 1, 768))
pos_embed = nn.Parameter(torch.zeros(1, 197, 768))
tokens = torch.cat([cls_token, tokens], dim=1) + pos_embed

# すべてのトークンが他のすべてのトークンに注意を向ける、事前正規化方式のエンコーダーブロック1つ
block = nn.TransformerEncoderLayer(
    d_model=768, nhead=12, dim_feedforward=3072, activation="gelu", batch_first=True, norm_first=True
)
print(block(tokens).shape)

torch.Size([1, 197, 768])を出力します。196個のパッチ・トークンとクラス・トークンはそれぞれ768次元のままで、ViT-Baseの12ブロックのうち次のブロックに渡せる状態です。

モデルのサイズとバリエーション#

ViTの名前はモデルサイズとパッチサイズを表します。ViT-B/16はパラメーター数8,600万のBaseモデルで、パッチサイズは16×16です。LargeモデルとHugeモデルのパラメーター数は、それぞれ3億700万と6億3,200万です。パッチを小さくするとトークン数が増え、精度が向上しますが、計算コストも増加します。ViTの学習方法や構造を変更したバリエーションがいくつかあります。

• DeiT： 強力なデータ拡張と蒸留トークンを使う学習レシピです。追加の事前学習データを使わず、ImageNet-1kだけで競争力のあるViTを学習できます。
• Swin Transformer： シフトした局所ウィンドウ内で注意を計算し、段階間でパッチを統合します。これによりCNNに似た階層構造が得られ、検出やセグメンテーションに適します。
• マスク付きオートエンコーダー（MAE）とDINOv2： ラベルなし画像を使ってViTを事前学習する自己教師あり学習手法です。MAEではマスクしたパッチを再構成し、DINOv2では教師ネットワークの出力を一致させます。
• CLIPとSigLIP： 画像とキャプションを対応付けるように学習したViT画像エンコーダーで、マルチモーダルモデルの視覚エンコーダーとして広く再利用されています。

用途#

• 視覚言語モデル： ViTエンコーダーは画像をトークン系列に変換します。これを大規模言語モデルがテキストとともに読み取り、画像キャプションの生成や視覚的質問応答を可能にします。
• 密な予測： ViTバックボーンのパッチ・トークンは、画像全体の広範囲からコンテキストを得る必要がある画像セグメンテーションや深度推定に活用されます。
• 動画理解： 動画ViTはクリップを時空間パッチに分割します。V-JEPAなどの自己教師あり動画モデルは、ラベルなし動画から動きや外観の特徴を学び、行動認識などのタスクに活用します。

Vision TransformerとCNNの比較#

CNNは、各フィルターが小さな近傍領域を見ることによる局所性と、同じフィルターがすべての位置に適用されることによる並進同変性という2つの前提を各層に組み込んでいます。パッチグリッドを除けば、ViTはデータから空間構造を学習する必要があります。元の論文では、ImageNetのみ（画像130万枚）で学習したViTは、同程度の規模のResNetより数パーセントポイント低いスコアを記録し、画像3億枚のJFT-300Mデータセットでの事前学習により、ViT-H/14のImageNetでのTop-1精度は88.55%に達しました。もう1つの違いは計算コストです。セルフアテンションでは、各トークンを他のすべてのトークンと比較するため、画像の一辺の長さを2倍にするとトークン数は4倍になり、アテンションの計算コストは16倍になります。フラッシュアテンションはこの計算のメモリ転送量を削減し、ウィンドウアテンションは比較そのものを制限します。UltralyticsはMetaのあらゆる対象をセグメント化するモデルを実行しており、そのsam_b.ptチェックポイントはViT-Bバックボーンを使用して各1024×1024画像をエンコードし、16×16のパッチが4,096個のトークンからなる64×64のグリッドを形成するため、12ブロックのうち8ブロックは14×14のウィンドウ内でのみアテンションを行い、残り4ブロックはグローバルアテンションを使用します。