【ITニュース解説】Vision Transformer (ViT) from Scratch in PyTorch

Vision Transformer（ViT）は、画像認識の分野に新たな風を吹き込んだ技術である。長年にわたり、画像認識のタスクではConvolutional Neural Network（CNN）が主流の座を占めてきたが、「An Image Is Worth 16x16 Words」という論文の発表以来、ViTはその強力な性能でCNNに挑戦し続けている。ViTの核心は、画像をテキスト（文章）のように扱うという斬新なアプローチにある。具体的には、画像を小さな「パッチ」に分割し、それらを「単語」の並びとして処理するのだ。この手法により、Transformerという元々自然言語処理で大きな成功を収めたモデルを、画像認識に応用することが可能になった。

今回の記事では、このVision TransformerをPyTorchという人気の機械学習フレームワークを使って、ゼロから実装し、簡単な画像分類のデータセットで学習させるプロセスが解説されている。対象となるのは、ピザ、ステーキ、寿司という身近な3種類の食べ物の画像だ。

ViTの基本的なアイデアは非常にシンプルである。まず、入力となる画像を固定されたサイズの小さなパッチ（例えば16×16ピクセル）に分割する。これは、ちょうど文章を個々の単語に分けるようなものだ。次に、これらのパッチをそれぞれ一次元のベクトルに「平坦化」する。これにより、各パッチはモデルが扱える「トークン」、つまりデータのかたまりとして表現される。

このトークンの並びには、いくつか特別な要素が追加される。一つは「[CLS]トークン」だ。これは、文章全体の意味を代表する特殊な単語のように、画像全体の情報を集約し、最終的な分類タスクに利用される。もう一つは「位置エンべディング」と呼ばれる情報である。画像をパッチに分割し、それらを単に並べるだけでは、元の画像におけるパッチの空間的な位置関係が失われてしまう。位置エンべディングは、各パッチが画像のどこにあったかという情報をモデルに教え込む役割を果たす。これにより、モデルはパッチの並びだけでなく、その配置からも意味を読み取ることができるようになる。これらの要素が加わったトークンのシーケンスは、Transformerエンコーダという強力なモジュールによって処理される。Transformerエンコーダは、各トークンが他のすべてのトークンとどのように関係しているか（これを「アテンション」と呼ぶ）を学習することで、画像全体の特徴を捉えるのだ。

記事では、ViTの基本構成である「ViT-Base」の設定例も紹介されている。例えば、入力画像は224×224ピクセルにリサイズされ、それを16×16ピクセルのパッチに分割すると、合計で196個のトークンが生成される。これらのトークンは、それぞれ768次元のベクトルとして表現される（埋め込み次元）。ViT-Baseは、12層のTransformerエンコーダを持ち、各層には12個のアテンションヘッドが組み込まれている。これらの設定全体で、モデルは約8580万個ものパラメータを持つ、非常に大規模なモデルとなる。

実際にモデルを学習させる際には、特定のデータセットが必要となる。今回の例では、ピザ、ステーキ、寿司の3つのカテゴリからなるデータセットが使われた。全ての画像は、モデルの入力要件に合わせて224×224ピクセルにリサイズされている。

学習のセットアップも重要だ。モデルのパラメータを最適化するためには、「Optimizer（最適化手法）」が必要であり、ここではAdamが用いられた。モデルの予測と実際の正解とのズレを測るための「Loss（損失関数）」には、分類タスクでよく使われるCrossEntropyLossが選ばれた。学習の進行度を制御する「学習率（LR）」は0.001に設定され、一度にモデルに与えるデータの量を示す「バッチサイズ」は32、データセット全体を何回学習させるかを示す「エポック数」は10に設定された。また、高速な計算を実現するために、GPU（CUDA）が利用された。

学習結果を見ると、いくつかの重要な点が明らかになる。まず、学習損失は非常に速いペースで減少していく。これは、ViTが非常に強力なモデルであることを示している。しかし、検証損失を見ると、初期段階は減少するものの、やがて横ばいになったり、時には上昇に転じたりする傾向が見られる。これは「過学習」と呼ばれる現象の兆候である。過学習とは、モデルが学習データにのみ最適化されすぎてしまい、まだ見たことのない新しいデータに対しては性能が落ちてしまう状態を指す。学習データの精度はほぼ100%に達する一方で、検証データの精度こそがモデルの真の汎用性能を反映していると言える。ViTのような大規模なモデルは、小さなデータセットで学習させると、すぐに過学習に陥りやすい特性がある。そのため、実際の応用においては、より大規模なデータセットで事前に学習された「事前学習済みViT」を利用し、それを自身のタスクに合わせて微調整する「ファインチューニング」という手法が非常に有効である。

この実験を通じて得られる重要な教訓は、アテンション機構が自然言語処理だけでなく、画像認識の分野でも非常に効果的であるというViTの証明だ。また、従来の画像認識がピクセル単位で画像を扱っていたのに対し、ViTは画像をパッチ、そしてトークンという抽象度の高い単位で捉えるという根本的なアプローチの変化を強調している。

今後のステップとして、より大規模なデータセット（例えばCIFAR-100やImageNetの一部）でViTを試すことや、HuggingFaceやtimmといったライブラリで提供されている事前学習済みモデルの重みを活用すること、さらにMixupやCutMixといったデータ拡張の手法を実験することが推奨されている。これらは、ViTの性能をさらに引き出し、過学習を抑制するための効果的なアプローチである。