【ViT】画像認識の大革命！Vision Transformerとは？

高橋

初めまして。AI-Tech-Lab愛知支部の学生エンジニアの高橋です。

初めてのブログということで、Vision Transformerのアーキテクチャについて簡単にまとめました。

Vision Transformerの概要

Vision Transformerとは、論文「An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale」で提案された画像認識の為のアーキテクチャです。

画像処理のDeep Learningと言えばCNNをベースとした畳み込みを用いたアーキテクチャが高いパフォーマンスを出していました。しかしながら、Vision Transformerは畳み込みをせずにそれを上回るほどの高いパフォーマンスを発揮したことにより注目を集めました。

Vision Transformerの特徴

・畳み込みを行わない

・SOTAを更新

・計算量を大幅に減少

Vision Transformerは畳み込みをせずにSOTAモデル以上のパフォーマンスを発揮しました。更にViT-H/14はCNNベースのBig Transfer (BiT)と比較してVision Transformerは計算量を約4分の1と大きく減少させました。このことにより今日(2022年)ではVision Transformerを元としたモデルの研究が数多く行われています。

以下ではViTのアーキテクチャをメインに解説していきます。

Vision Transformerのアーキテクチャ・処理の解説

Vision Transformerは主にTransformerのエンコーダ部分で構成されています。

Vision Transformerの処理の流れ

入力画像をベクトルに変換
Transformer Encoderにより処理
MLP Headにより処理

①入力画像をベクトルへ変換

まず入力画像 $\boldsymbol{x} \in \mathbb{R}^{H \times W \times C}$ をパッチに分割します。
ここで(H, W)は解像度、Cはチャンネル数を表しています。

次に $\boldsymbol{x}_{p} \in \mathbb{R}^{N \times (P^2 \cdot C)}$ というベクトルに変換します。
ここでPはパッチの大きさ、Nはパッチ数 $N=HW/P^{2}$ を表しています。

更に学習可能な $\boldsymbol{E} \in \mathbb{R}^{(P^2 \cdot C) \times D}$ と位置エンコーディング $\boldsymbol{E}_{pos} \in \mathbb{R}^{(N+1) \times D$ を用いて式(1)の様にマッピングします。

(1) $\begin{equation*}\boldsymbol{z_{0}} = [\boldsymbol{x}_{class};\: \boldsymbol{x}_{p}^1\boldsymbol{E};\: \boldsymbol{x}_{p}^2\boldsymbol{E};\: \cdots; \boldsymbol{x}_{p}^N\boldsymbol{E}] + \boldsymbol{E}_{pos}\end{equation*}$

以上の処理がTransformerのEncoder前までの処理となります。

②Transformer Encoderにより処理

MSAをMultihead Self-Attention、MLPをMulti Layer Perceptron、LNをLinear NormlizationとするとTransformerのEncoderの処理は式(2)と式(3)で表すことができます。
MLPの活性化関数はGELUを用いています。

(2) $\begin{equation*}z'_l = MSA(LN(z_{l-1})) + z_{l-1} \hspace{1pc} l=1 \cdots L\end{equation*}$

(3) $\begin{equation*}z_l = MLP(LN(z'_l)) + z'_l \hspace{1pc} l=1 \cdots L\end{equation*}$

③MLP Headにより処理

$Z^0_L$ は最終層の出力における前から0番目のベクトルを表しています。つまり、[cls]トークンの最終出力です。それをNLに入れて得たyをMLP Headに入れることで最終予測を得ることができます。

(4) $\begin{equation*}y=LN(z^0_L)\end{equation*}$

Multi-Head Attentionとは

Multi-Head AttentionはTransformerのencoderで使われている機構になります。この機構の優れている点はどの情報に注目すべきか自動で判断をして情報を処理する点です。

Scaled Dot-Product Attention

Multi-Head Attentionの説明の前にそれの元となるScaled Dot-Product Attentionについて解説をします。
Scaled Dot-Product Attentionは式(5)で表すことができます。

(5) $\begin{equation*}Attention(Q,K,V)=softmax(\frac{QK^t}{\sqrt{d}})V\end{equation*}$

この式は引数としてQ(query)、K(key)、V(value)を取ります。この式はバッチ処理を行う式であるので、まずは入力が一つの式について解説を行います。入力が一つの場合の式(6)はとなります。

(6) $\begin{equation*}Attention(\boldsymbol{q},K,V)=softmax(\frac{\boldsymbol{q}K^t}{\sqrt{d}})V\end{equation*}$

まず $\boldsymbol{q}K^t$ を計算すると式(7)となり、内積が並ぶことになります。
内積は類似度を表現しているので、結果はqueryとkeyの類似度を並べたベクトルとなります。

(7) $\begin{equation*}\boldsymbol{q}K^t &=& \boldsymbol{q}(\boldsymbol{k}_1^t, \boldsymbol{k}_2^t, \cdots ,\boldsymbol{k}_m^t) \\&=& (\boldsymbol{q} \cdot \boldsymbol{k_1}, \boldsymbol{q} \cdot \boldsymbol{k_2}, \cdots, \boldsymbol{q} \cdot \boldsymbol{k_m})\end{equation*}$

次に勾配消失を抑制する為 $\sqrt{d}$ で割って補正を行います。

更にそれをソフトマックス関数に入れることにより類似度を合計の重みに変換し、それをPとします。

(8) $\begin{equation*}softmax(\frac{\boldsymbol{q}K^t}{\sqrt{d}})=(p_1,p_2,\cdots,p_m)=\boldsymbol{P}\end{equation*}$

最後にVをかけると

(9) $\begin{equation*}softmax(\frac{\boldsymbol{q}K^t}{\sqrt{d}})V=\boldsymbol{P}V=\sum p_i\boldsymbol{v_i}\end{equation*}$

従って、Attention(q,K,V)ではqueryとkeyの類似度を計算し、それに応じた重みのvalueを出力することになります。
以上の過程を式にすると式(10)、バッチ処理を行う式は式(11)となります。

(10) $\begin{equation*}Attention(\boldsymbol{q}, K, V)=softmax(\frac{\boldsymbol{q}K^t}{\sqrt{d}})V=\boldsymbol{P}V=\sum p_i\boldsymbol{v_i}\end{equation*}$

(11) $\begin{equation*}Attention(Q, K, V)=softmax(\frac{QK^t}{\sqrt{d}})V= \left( \begin{array}{c}p_{11}\boldsymbol{v}_{1}+ \cdots + p_{1m}\boldsymbol{v}_{n} \\\vdots \\p_{m1}\boldsymbol{v}_{1}+ \cdots + p_{nm}\boldsymbol{v}_{n} \end{array} \right)\end{equation*}$

Multi-Head Attention

次にVision Transformerで使われているMulti-Head Attentionについて解説をします。
Multi-Head Attentionは式(12)となります。

(12) $\begin{gather*}Multi-Head(Q, K, V)=concat(head_{i})W^{o} \\head_{i} = Attention(QW^{Q}_{i}, KW^{K}_{i}, VW^{V}_{i})\end{gather*}$

Multi-Head AttentionはScaled Dot-Product Attentionを主として構成されています。
Q、K、Vそれぞれに学習可能な行列をかけてAttentionに入力し、その結果を連結し、更に行列をかけたものがMulti-Head Attentionとなります。

様々な角度から入力ベクトルを比較し、どこに注目すべきか判断をして、出力を決定します。

Attentionについて詳しく知りたい方はAttention Is All You Needをご覧下さい。

まとめ

今回はVision Transformerのアーキテクチャをメインにまとめてみました。Vision TransformerはTransformerのエンコーダを用いて高いパフォーマンスを発揮しています。より詳細にVision Transformerについて知りたい方は原文を参考にして下さい。

今日(2022年)ではSwin TransformerやData-efficient Image TransformerなどVision Transformerを元としたモデルが多く研究されており、日々の勉強が欠かせません。

私たち学生エンジニア集団「AI-Tech-Lab」では、今回紹介したVision Transformerなどの画像処理だけでなく機械学習に関わるものを幅広く学んだり、PoC開発を行ったりしています。

▼AI-Tech-Labについての詳細はこちら