ViTとは?画像をパッチとして読むTransformer

画像をパッチ列へ変換してTransformer Encoderで分類するVision Transformerの全体構造

画像をパッチ列へ変換してTransformer Encoderで分類するVision Transformerの全体構造

Vision Transformer(ViT)は、画像を小さなpatch(固定サイズの領域)へ分け、各patchをtokenのように並べてTransformer Encoderで処理する画像認識modelです。CNN(Convolutional Neural Network:畳み込みニューラルネットワーク)を主役にせず、標準的なTransformerへ最小限の変更を加えて画像分類へ適用しました。

「画像を単語として読む」という説明は直感的ですが、pixelを文字へ変換するわけではありません。patchをvectorへ投影し、位置情報を加えたsequenceとしてSelf-Attentionへ渡します。どのshapeがどう変わるかを追うと、ViTの仕組みと計算costが見えてきます。

3行で分かるVision Transformer

  • 画像を\(P\times P\)のpatchへ分け、各patchを\(D\)次元のPatch Embeddingへ変換します。
  • 先頭へ学習可能なClass Tokenを加え、Position Embeddingと共にTransformer Encoderへ入力します。
  • Self-Attentionで全patch間の関係を学び、最後のClass Tokenから画像classを予測します。

この記事では、ViT原論文Google Researchの公式repositoryをもとに、224×224画像をViT-B/16へ入力する具体例、Encoder内部、CNNとの違い、patch sizeと解像度のtrade-offまで説明します。

224×224画像は197個のtokenになる

最初に、RGB画像1枚をViT-B/16へ入れる例を最後まで計算します。

  • 画像の高さ\(H=224\)
  • 画像の幅\(W=224\)
  • channel数\(C=3\)
  • patch size \(P=16\)
  • embedding dimension \(D=768\)

縦横を16 pixelずつに分けると、patch gridは\(14\times14\)です。

\[\frac{224}{16}\times\frac{224}{16}=14\times14=196\]

1 patchは\(16\times16\times3=768\)個のpixel値を持ちます。これをflatten(多次元配列を1本のvectorへ並べる処理)すると、画像全体は196本の768次元vectorになります。

さらにsequence先頭へClass Tokenを1個追加するため、Transformerへ入るtoken数は197です。

\[(224,224,3) \rightarrow(196,16,16,3) \rightarrow(196,768) \rightarrow(197,768)\]

224×224画像を16×16パッチへ分割し197個のViT tokenへ変換する流れ

この例で「patch内のpixel数」と「embedding dimension」がどちらも768なのはViT-B/16固有の一致です。別のpatch sizeやmodel sizeでは同じとは限りません。

Patch Embeddingは画像領域をtoken表現へ写す

入力画像を\(x\in\mathbb{R}^{H\times W\times C}\)とします。\(P\times P\)へ分割したpatch数\(N\)は次のとおりです。

\[N=\frac{H}{P}\frac{W}{P}=\frac{HW}{P^2}\]

各patchをflattenすると、patch列\(x_p\)のshapeは\(N\times(P^2C)\)になります。学習可能な行列\(E\in\mathbb{R}^{(P^2C)\times D}\)を掛けると、各patchを\(D\)次元へ投影できます。

\[x_pE\in\mathbb{R}^{N\times D}\]

このprojectionがPatch Embeddingです。tokenizerが単語境界を選ぶ自然言語処理とは違い、original ViTは内容に関係なく規則正しいgridで画像を切ります。patchは意味的な物体単位ではなく、modelが処理する計算単位です。

flatten+LinearはConv2dでも実装できる

patchを1枚ずつ切り出してflattenし、同じLinear layerへ通す処理は、kernel sizeとstrideをどちらも\(P\)にしたConv2d(2次元畳み込み)でも実装できます。

\[\operatorname{Conv2d}(C,D,\text{kernel}=P,\text{stride}=P)\]

flattenとLinearによるPatch EmbeddingをConv2dで等価に実装する条件

kernelがpatch全体を一度に覆い、strideが同じ幅だけ進むため、patch同士は重なりません。weightを対応する順序へ並べれば、各patchに同じ線形変換を適用する点でflatten+Linearと等価です。

ただし、これは効率的な実装方法の等価性です。「ViTがCNNと同じ多層畳み込みで特徴を階層化する」という意味ではありません。projection後の主な処理はTransformer Encoderです。

Class TokenとPosition Embeddingは役割が違う

Patch Embeddingだけでは\(N\)本の画像tokenがあります。ViTはその先頭へ、学習可能なvector\(x_{\mathrm{class}}\in\mathbb{R}^{D}\)を1本追加します。これがClass Tokenです。

Class Tokenは正解class IDを入力するtokenではありません。すべての画像で同じ初期vectorから始まり、Self-Attentionを通じてpatch情報を集約します。最終層のClass Tokenが画像全体の表現としてclassification headへ渡されます。

sequenceの初期状態\(z_0\)は次の式です。

\[z_0= \left[x_{\mathrm{class}}; x_p^1E; x_p^2E; \ldots; x_p^NE\right]+E_{\mathrm{pos}}\]

Class TokenとPatch TokenへPosition Embeddingを加えるViT入力sequence

\(E_{\mathrm{pos}}\in\mathbb{R}^{(N+1)\times D}\)が学習可能なPosition Embeddingです。Class Tokenを含む各位置へ別のvectorを足します。

Self-Attentionは入力順を入れ替えると出力も同じように入れ替わる計算です。位置情報がなければ、左上のpatchと右下のpatchを内容だけで区別できません。Position Embeddingによって、同じ見た目のpatchでも「画像gridのどこにあるか」を学べます。

original ViTが入力へ加えるのは学習可能な1D Position Embeddingです。ただし高解像度へ移すときは、patch token部分を2D gridとして補間します。入力形式はsequenceでも、元の二次元配置を完全に無視しているわけではありません。

Transformer Encoderは全patch間の関係を更新する

ViTのEncoder blockは、Multi-Head Self-Attention(複数のAttention headでtoken間関係を計算する層)とMLP(多層パーセプトロン)で構成されます。各sub-layerにはLayer Normalizationとresidual connection(入力を出力へ加える経路)があります。

原論文の式はPre-LN、つまりAttentionやMLPの前にLayer Normalizationを置く構成です。layer \(\ell\)の計算は次のようになります。

\[z’_{\ell} =\operatorname{MSA}(\operatorname{LN}(z_{\ell-1}))+z_{\ell-1}\]
\[z_{\ell} =\operatorname{MLP}(\operatorname{LN}(z’_{\ell}))+z’_{\ell}\]

LayerNorm、Multi-Head Self-Attention、MLP、残差接続からなるViT Encoder block

MLPは2層のLinear layerとGELU(滑らかな非線形活性化関数)を使います。ViT-Bではこのblockを12回繰り返しますが、token数とembedding dimensionはblock前後で\(197\times768\)のままです。

Self-Attentionでは各tokenからQuery、Key、Valueを作ります。1 headの基本計算は次のとおりです。

\[\operatorname{Attention}(Q,K,V) =\operatorname{softmax}\left(\frac{QK^{\mathsf T}}{\sqrt{d_k}}\right)V\]

画像グリッド上の一つのパッチが全パッチを参照するViTのSelf-Attention

\(QK^{\mathsf T}\)はtoken対tokenのscore matrixです。patch tokenは隣接patchだけでなく、離れたpatchやClass Tokenも直接参照できます。Class Token側も全patchを参照しながら層ごとに更新されます。

ただし「全patchを直接参照できる」は、すべてのheadが常に画像全体を均等に見るという意味ではありません。Attention weightは入力と学習結果で変わり、headごとに異なる関係を表します。

Transformerそのものの構造を先に確認したい場合は、Transformerとは?Attention Is All You Need論文からLLMの基本構造を解説も参照してください。

最後のClass Tokenから画像classを予測する

\(L\)個のEncoder blockを通った後、先頭tokenだけを取り出してLayer Normalizationします。

\[y=\operatorname{LN}(z_L^0)\]

\(z_L^0\)は最終layerの0番目、つまりClass Tokenです。fine-tuningではこれを1層のLinear headへ入力し、classごとのlogitを計算します。

patch tokenの最終出力が消えるわけではありません。original ViTの画像分類ではClass Tokenをheadへ使う、という選択です。物体検出やsegmentationのように位置ごとの予測が必要なtaskでは、patch featureを利用する別のheadやarchitectureが必要です。

ViT-B/16のBと16は何を表すのか

ViT-B/16は、Base sizeのTransformerを使い、patch sizeが16×16であることを表します。BLHはEncoderの深さと幅、/16/32はpatchの細かさです。

Model Encoder layers Hidden size \(D\) MLP size Heads Parameters
ViT-Base 12 768 3,072 12 約86M
ViT-Large 24 1,024 4,096 16 約307M
ViT-Huge 32 1,280 5,120 16 約632M

同じBaseでもViT-B/32ViT-B/16ではtoken数が違います。同じpatch sizeでもBaseとLargeでは各tokenを処理するnetwork規模が違います。model sizeとpatch sizeは別々の計算cost要因です。

patchを細かくするとdetailと計算costが増える

224×224画像でpatch sizeだけを変えてみます。Class Tokenを含むtoken数\(T=N+1\)と、1 headあたりのAttention score matrix要素数\(T^2\)を比較します。

Patch size Grid Tokens \(T\) Attention要素数 \(T^2\) 特徴
32×32 7×7 50 2,500 粗い領域、計算は軽い
16×16 14×14 197 38,809 detailとcostの中間
8×8 28×28 785 616,225 細かい領域、計算は重い

ViTのパッチサイズ32、16、8におけるトークン数とAttention計算量の比較

patchの一辺を16から8へ半分にすると、patch数は約4倍です。Attention matrixはtoken数の二乗なので、要素数は約15.9倍になります。

\[\frac{785^2}{197^2}\approx15.9\]

小さいpatchは細い線や小物体を早い段階で失いにくい一方、memoryと計算量を大きく増やします。逆に大きいpatchは軽量ですが、patch内の細かな配置を最初の1 vectorへ圧縮します。

「小さいpatchほど必ず高精度」とは限りません。model規模、training data、augmentation、解像度、対象物の大きさを含めて評価する必要があります。

解像度を上げると位置埋め込みの長さが変わる

224×224、P=16で事前学習したmodelを384×384へfine-tuningするとします。patch sizeを変えなければgridは14×14から24×24へ増えます。

\[224:\quad14\times14+1=197\ \text{tokens}\]
\[384:\quad24\times24+1=577\ \text{tokens}\]

224から384への高解像度fine-tuningで位置埋め込みを二次元補間する流れ

学習済みPosition Embeddingは197本しかないため、そのまま577 tokensへ足せません。原論文ではClass Tokenの位置を分け、14×14のpatch positionを二次元gridとして24×24へ補間します。

この方法は、学習済みの空間配置を新しい解像度へ連続的に伸ばす近似です。新しい位置を観測済みにするわけではないため、補間後のdataでfine-tuningして調整します。

解像度を224から384へ上げるとtoken数は約2.93倍、Attention要素数は約8.58倍になります。patch sizeを固定した高解像度化も大きなcost増です。

ViTとCNNは何が違うのか

CNNは小さなkernelを画像上で共有し、近傍のpatternを段階的に組み合わせます。locality(近いpixelほど関係が強いという仮定)とtranslation equivariance(入力位置がずれるとfeature位置も対応してずれる性質)をarchitectureへ強く組み込みます。

original ViTは画像をpatchへ分けた後、Transformer内部で全tokenを同じ形式として扱います。近傍だけを見る制約や二次元の平行移動構造はCNNほど強くありません。このようにmodelへあらかじめ組み込む仮定をinductive bias(学習前から持つ問題固有の偏り)と呼びます。

観点 CNN original ViT
初期単位 局所pixel領域 固定size patch
情報交換 local kernelを重ねて範囲を広げる Self-Attentionで全token間を直接接続できる
位置構造 局所性とweight共有を強く組み込む 学習可能なPosition Embeddingから関係を学ぶ
dataが少ない場合 画像向けbiasが助けになることがある 規則をdataから学ぶ負担が大きい
主なcost feature map、kernel、channelに依存 token数の二乗に依存するAttentionが重い
長所 局所patternを効率よく再利用 離れたpatch間を早い層から結び付けられる

CNNの局所畳み込みとViTのグローバルSelf-Attentionにおける帰納バイアスの比較

これは「CNNはglobal情報を扱えない」「ViTにはlocalityがない」という二択ではありません。CNNも層を重ねれば広いreceptive field(出力へ影響する入力範囲)を持ちます。ViTもpatch化と学習後のAttention patternを通じて局所的な関係を表せます。違いは、どの構造を最初から強く仮定し、どこをdataから学ぶかです。

原論文の結果は大規模事前学習とセットで読む

ViT原論文はImageNet約1.3M枚、ImageNet-21k約14M枚、JFT-300M約303M枚を使い、data規模を変えて事前学習しました。

小さいImageNetだけで学習した条件では、大きなViTが同等規模のResNetより低い結果になる場合がありました。一方、ImageNet-21kやJFT-300Mで事前学習したViTは、複数の画像分類benchmarkで強い結果を示しました。

したがって原論文の中心的な発見は、「Transformerへ画像を入れればdata量に関係なくCNNを超える」ではありません。画像固有のinductive biasが弱いmodelでも、十分なdataと計算で事前学習すれば高い転移性能を得られる、という条件付きの結果です。

現在のmodelを選ぶときも、原論文のmodel名だけで決めず、使うcheckpointの事前学習data、license、入力解像度、fine-tuning条件を確認します。公式ViT repositoryはactiveな研究実装の後継としてbig_visionを案内しています。

ViTが苦手になりやすい条件

小さな物体や細い構造

最初のpatch化で領域を1 vectorへまとめるため、patchより十分小さい特徴は表現しにくくなることがあります。小さいpatchや高解像度は対策候補ですが、token数とcostが増えます。

長いtoken列

標準Self-Attentionのscore matrixは\(T\times T\)です。高解像度画像、動画、複数画像ではtoken数が急増します。memory不足やlatency制約がある場合は、patch size、解像度、token削減、階層型architectureなどを検討します。

位置を直接出すtask

Class Tokenからclassを出すoriginal ViTは画像分類modelです。どのpatchに物体があるか、pixelごとのclassは何かを、そのheadだけでは出力しません。検出やsegmentationにはpatch featureを利用するtask-specific headと学習が必要です。

Attention mapの過剰解釈

Attention weightを可視化すると参照関係の一面を観察できますが、それだけで予測の因果的な理由を証明したことにはなりません。複数head、MLP、residual経路を通るため、説明手段と忠実な因果説明を分けて扱います。

ViTを設計・利用するときの確認項目

判断軸 確認すること trade-off
Patch size 最小対象物や必要detail 小さいほどtokenとcostが増える
Input resolution 撮影条件と細部 高いほど位置埋め込みと計算量が増える
Model size latency、memory、精度 B/L/Hで深さと幅が大きく変わる
Pre-training data domain、規模、license transfer先とのずれが性能へ影響する
Task head 分類、検出、segmentation Class Tokenだけで位置出力はできない
Evaluation class別、サイズ別、domain別 平均accuracyだけでは弱点が隠れる

CLIPではViTをImage Encoderとして使い、最終画像表現をtext embeddingと比較できます。ViTが画像をtoken列へ変換する部分と、CLIPが画像とtextを同じ空間へそろえる部分は別の責務です。接続関係はCLIPとは?画像とテキストを同じ空間で扱う技術で解説しています。

まとめ

Vision Transformerは、画像を固定sizeのpatchへ分割し、Patch Embedding、Class Token、Position Embeddingを組み合わせてTransformer Encoderへ入力します。ViT-B/16で224×224画像を扱う場合は196 patches、Class Token込み197 tokensです。

設計上の重要点は、patchをtokenとして扱えることだけではありません。Self-Attentionで離れた領域を直接結び付ける一方、patchを細かくしたり解像度を上げたりするとAttention costが急増します。またCNNより画像固有のinductive biasが弱いため、原論文の強い結果は大規模事前学習という条件と一緒に読む必要があります。

次回は、ViTなどのVision EncoderをLLMへ接続するVLM(Vision-Language Model:画像と言語を扱うmodel)を取り上げます。画像tokenを作った後、言語modelが理解できる表現へどう橋渡しするのかを整理します。

参考文献

コメント

タイトルとURLをコピーしました