MLエンジニアへの道 #30 - ViT

MLDLCV

Last Edited: 11/20/2024

このブログ記事では、ディープラーニングにおけるViTについて紹介します。

ML

これまで、注意機構やトランスフォーマーについては、主に自然言語処理や類似する時系列データを扱うタスクにおける文脈で議論してきました。 しかし、これらはコンピュータビジョンのタスクにも応用可能です。

画像の解釈

画像処理の課題の1つは、その膨大なデータ量です。一般的なコンピュータビジョンで扱う画像のサイズは224×224ピクセルで、 カラー画像の場合は約15万個のピクセル値になります。これらの個々のピクセルを入力単位として処理することは、 トランスフォーマーでは現実的ではありません。畳み込みカーネルを共有するCNN（畳み込みニューラルネットワーク）が、 トランスフォーマーが自然言語処理（NLP）で成功を収めた後も、依然としてコンピュータビジョンの分野で主流となっていたのには、 そういった背景があります。

class Patches(layers.Layer):
    def __init__(self, patch_size, **kwargs):
        super(Patches, self).__init__(**kwargs)
        self.patch_size = patch_size
 
    def call(self, images):
        batch_size = tf.shape(images)[0]
        patches = tf.image.extract_patches(
            images=images,
            sizes=[1, self.patch_size, self.patch_size, 1],
            strides=[1, self.patch_size, self.patch_size, 1],
            rates=[1, 1, 1, 1],
            padding="VALID",
        )
        patch_dims = patches.shape[-1]
        patches = tf.reshape(patches, [batch_size, -1, patch_dims]) # batch_size, patch_num, patch_dims
        return patches

しかし、Dosovitskiyら（2021）は、画像を16×16ピクセル程度のパッチに分割し、それらを単語として扱うことで、 CNNを使わずにトランスフォーマーを画像に応用できることを提案しました。この手法により、 入力単位をトランスフォーマーが処理可能な量（15万ピクセルから588パッチ）まで減らせます。 線形層を使用して各パッチの次元を768（16×16×3）からさらに削減し、可能な限り情報を保持した状態で適切なサイズのパッチ埋め込みを作成できます。

学習可能な位置埋め込み

しかし、画像をパッチに分割すると、位置情報が失われるという問題があります。NLPタスクのようにサイン波やコサイン波を用いた位置エンコーディングは、 文のような1次元の順序データに適している一方で、2次元画像の空間的な位置を直接表現することはできません。

Positional Embedding

この問題に対処するため、Dosovitskiyら（2021）は、パッチ埋め込みと同じサイズの学習可能な位置埋め込みを導入しました。 これらの埋め込みは、2次元画像内のパッチ間の空間的関係をキャプチャするように学習されます。学習後、これらの位置埋め込みは、 2次元空間内のパッチの適切な位置を表現することに成功しており、その有効性が確認されています。これにより、 パッチ埋め込みと位置埋め込みを準備したうえで、それらをトランスフォーマーに渡してコンピュータビジョンタスクを実行できます。

ビジョントランスフォーマー

画像分類のための注目すべきトランスフォーマーアーキテクチャの1つが、ビジョントランスフォーマー（Vision Transformer, ViT）です。 これは、パッチ埋め込み、位置埋め込み、追加の学習可能なクラス埋め込みを処理するトランスフォーマーエンコーダーで構成されています。 このクラス埋め込みは、画像分類を学習する際に使用され、クラス情報に対応する出力を生成します。以下にViTアーキテクチャの概要を示します。

ViT

以下は学習可能な位置埋め込み、追加のクラストークン、およびViTモデルのTensorFlowでの実装です。トランスフォーマーエンコーダーが何か分からない場合は、 前回の記事MLエンジニアへの道 #29 - BERT vs BPTを参照してください。

class PatchClassEmbedding(layers.Layer):
    def __init__(self, d_model, n_patches, kernel_initializer="he_normal", **kwargs):
        super(PatchClassEmbedding, self).__init__(**kwargs)
        self.n_tot_patches = n_patches + 1
        self.kernel_initializer = tf.keras.initializers.get(kernel_initializer)
        self.class_embed = self.add_weight(
            shape=(1, 1, d_model),
            initializer=self.kernel_initializer,
            name="class_token",
        )  # extra learnable class
        self.position_embedding = layers.Embedding(
            input_dim=(self.n_tot_patches), output_dim=d_model
        )
 
    def call(self, inputs):
        positions = tf.range(start=0, limit=self.n_tot_patches, delta=1)
        x = tf.repeat(self.class_embed, tf.shape(inputs)[0], axis=0)
        x = tf.concat((x, inputs), axis=1)
        encoded = x + self.position_embedding(positions)
        return encoded
 
class ViT(tf.keras.Model):
  def __init__(self, num_classes, patch_size, num_patches, embed_dim, 
               hidden_dim, num_layers=4, num_heads=8):
    super(ViT, self).__init__()
    self.patches = Patches(patch_size)
    self.linear = layers.Dense(embed_dim)
    self.embedding = PatchClassEmbedding(embed_dim, num_patches)
    self.encoder = tf.keras.Sequential([
        TransformerEncoderBlock(embed_dim, hidden_dim) 
        for i in range(num_layers)
    ])
    self.classifier = tf.keras.Sequential([
        layers.Lambda(lambda x: x[:,0,:]), ## take the class embedding for classification
        tf.keras.layers.Dense(num_classes, activation='softmax')
    ])
 
  def call(self, x):
    x = self.patches(x)
    x = self.linear(x)
    x = self.embedding(x)
    x = self.encoder(x)
    x = self.classifier(x)
    return x

ViTのアーキテクチャはBERTと非常に類似しており、データ前処理のフェーズが主な違いです。しかしViTは、大規模なデータセットで事前学習を行い、 ファインチューニングすることで、小規模から中規模の画像認識タスクで最先端（SOTA）の結果を達成しています。この成功の背景には、 トランスフォーマーの持つ少ない帰納的バイアスが関与していると考えられます。これにより、十分なデータが提供されれば、 より柔軟に複雑なパターンを学習できるようになります。

結論

前回の記事と同様に、あえてパイプラインの実装やPyTorchの例を省略しています。実際に自分で試してみてください。 別のアプローチとして、前処理フェーズを畳み込み層に置き換えたハイブリッドビジョントランスフォーマーを構築することも可能です。 この手法は、小規模なモデルに対して特に効果的であることが示されています（大規模モデルでは、ViTとの性能差はほとんどありません）。 ハイブリッドビジョントランスフォーマーの実装にも挑戦してみてください。

ViT（Vision Transformer）は、トランスフォーマーのアーキテクチャが持つ柔軟性を強調しており、現在ではNLPを超えてさまざまなデータタイプやタスクに使用されています。 本記事では多くの詳細や知見を割愛していますので、以下で引用したオリジナルの論文をぜひお読みいただくことをお勧めします。

リソース

TOPに戻る↑