More than 1 year has passed since last update.

Vision Transformerはどのように働くのか？

Posted at 2023-07-01

原論文

HOW DO VISION TRANSFORMERS WORK?
https://arxiv.org/pdf/2202.06709.pdf

要約

multi-head self attentionは小さいデータセットで学習させたときに過学習は起こらない
multi-head self attentionは精度を向上させるだけでなく，convolution nerual networkよりも損失関数が平らになる
multi-head self attentionを使うと損失関数の凹凸が多くなってしまうが，大規模データセットと損失関数の平滑化手法によって凹凸を緩和することができる
multi-head self attentionはローパスフィルタの役割を持っている
convolution neural networkはハイパスフィルタの役割を持っている

multi-head self attentionの特徴

Q=XW_{Q}\\
K=XW_{K}\\
V=XW_{V}\\
z_{j}=\sum_{i} Softmax(\frac{QK}{\sqrt{d}})_{i} V_{i,j}

$X, Q, K, V$はそれぞれ入力する特徴，query，key，valueを表す．$d$はqueryやkeyのトークンの次元数を表す．$i$は$i$番目のヘッドを表す．$j$は$j$番目のトークンであることを表す．
$Softmax(\frac{QK}{\sqrt{d}})$は$V$に対する大きいカーネルであると考えることができる．これより，遠い距離にある特徴から特徴抽出できるようになった
$Softmax(\frac{QK}{\sqrt{d}})$は入力する特徴$X$によってカーネルの重みが動的に変わると考えることができる．
入力に近い層ではconvolution neural networkのように近傍の特徴から特徴抽出することが知られている
出力に近い層では遠くの特徴から特徴抽出することが知られている

convolution neural networkの特徴

Y_{i,j}=\sum_{m}\sum_{n}W_{m, n}X_{i+m, j+n}+b

$W, b, X, Y$はそれぞれ重み，バイアス，入力特徴，出力特徴を表す．$i，j$は特徴の$i$行$j$列であることを表す．$m，n$は重みの$m$行$n$列であることを表す．
カーネルサイズは$3\times3$や$5\times5$のことが多く，小さいカーネルである．これより，近傍の特徴からしか特徴抽出できない
$W$は学習によって更新されるが，入力$X$に対して変わることがない．つまり，カーネルの重みは入力$X$に対して動的に変わらない．
Dilated Convolutionを使うことで，通常のconvolution neural networkと比べて遠くの特徴から特徴抽出できるようになる．

vision transformerは小さいサイズのデータセットを学習させたときに過学習は起こらない

これまでは，小さいサイズのデータセットを使ってvision transformerを学習させると過学習が起こると考えられていた．しかし，実験によってvision transformerは過学習が起きていないことが分かった．

上図はCIFAR-100を用いたクラス分類の結果である．$Error$はテストデータを誤識別した割合，$NLL_{train}$は学習時の損失関数の値である．$RX$はResNeXt，$R$はResNet，$Twins$，$Swin$，$PiT$，$ViT$はそれぞれmulti-self attentionを採用したモデルである．
図(a)を見ると，テストデータを誤識別した割合と学習時の損失関数の値は比例関係のようになっていることが分かる．過学習とは，学習時のデータの損失が小さくなるが，テストデータの誤識別は大きくなることを指すため，図(a)の結果からmulti-head self attentionのモデルは過学習を起こしているとは言えない．multi-head self attentionは学習時の損失関数の値が大きいため画像の特徴をうまく学習できていないと言える．
また，図(b)を見ると，学習時のデータの数が増えるほど，テストデータを誤識別した割合と学習時の損失関数の値は小さくなっていく．これは学習データの数が増えるほど，multi-head self attentionは画像の特徴をうまく学習できていることを示している．

vision transformerの損失関数は凹凸が多くなる

小さいサイズのデータセットを学習させるときに，vision transformerが上手く学習できないのはなぜなのだろうか．その答えとして

convolution neural networkを用いた時と比べて，損失関数が平らになりやすい
convolution neural networkを用いた時と比べて，損失関数の凹凸が多くなる

の2つが挙げられる．

図(b)はヘッシアン最大固有値のスペクトラムを表す．横軸はヘッシアン固有値，縦軸はそのヘッシアン固有値の多さを表す．

ここで，ヘッセ行列とヘッシアン固有値について簡単に説明する．
ヘッセ行列とは，ある多変数関数が与えられたときに，その多変数関数を2階偏微分し，その偏微分値を行列にしたものである．深層学習における損失関数は数多くの重みをもっているため，多変数関数と考えることができる．
ヘッシアン固有値とは，ヘッセ行列の固有値を計算したものである．ヘッシアン固有値がすべて0より大きければ，狭義凸関数となっていることが分かる．ヘッシアン固有値がすべて0以上であれば，凸関数となっていることがわかる．ヘッシアン固有値に負の値があれば，凹凸のある関数となっていることが分かる．深層学習における損失関数は負に近づいていけば近づいていくほどよいことが多い．そのため，損失関数は下に凸の関数であることが望ましい．また，ヘッシアン固有値の絶対が大きいと関数の傾きが大きいことを示し，ヘッシアン固有値の絶対が小さいと関数の傾きが小さいことを示す．

ここで図(b)をもう一度確認する．
図(b)の左側は学習開始直後のヘッシアン最大固有値の分布を示している．$ViT$や$PiT$，$Swin$といったmulti-head self attentionのモデルは損失関数のヘッシアン最大固有値が0近くに多く分布していることがわかる．また，負の固有値にも多く分布していることが分かる．また，$ResNet$といったconvolution neural networkのモデルは損失関数のヘッシアン最大固有値が0より大きいところに広く分布しており，負の固有値を持たないことも分かる．これより，multi-head self attentionのモデルは損失関数が平らになりやすく，損失関数の凹凸が多くなることが分かる．
図(b)の右側は学習開始して300回学習したときのヘッシアン最大固有値の分布を示している．multi-head self attentionのモデルは全体的に損失関数のヘッシアン最大固有値が0より大きいところに分布するようになっていることが分かる．
図(a)は損失関数を可視化したものである．$ResNet$が損失関数の傾きが一番大きく，次に$Swin$が損失関数の傾きが大きくなっている．$ViT$や$PiT$は損失関数の傾きが平らに近くなっている．

深層学習は損失関数の微分値を使って損失関数が最小になるように学習する．このため，損失関数の微分値が大きいとその分だけ重みを大きく更新することができるようになる．また，凹凸が多いと損失関数の最小値を探すことが大変になるため，凹凸はなければないほうが良い．

vision transformerの損失関数の凹凸を緩和するには

データセットを大量に用意する
(クラス分類タスクにおいて)後段のヘッドではCLS tokenの代わりにglobal average poolingを使う
Swin Transformerのようにmulti-head self attentionを行う領域を制限する

上図は$ResNet$と$ViT$のヘッシアン最大固有値の分布を示している．また，$ViT$(6%)は学習データの6%を使って学習したことを示す．$ViT$と$ViT$(6%)を比べて分かるように，$ViT$(6%)は負の固有値を多くとる．これより，少ないデータセットでは損失関数の凹凸が多くなってしまうことが分かる．これより，vision transformerの損失関数の凹凸を緩和するには大量のデータセットを用意すれば良いことが分かる．

上図はクラス分類タスクにおいて，CLS tokenを使ってクラス分類したときの損失関数のヘッシアン最大固有値の分布とglobal average poolingを使ってクラス分類したときの損失関数のヘッシアン最大固有値の分布を示している．CLS tokenを使うよりもglobal average poolingを使った方が負の固有値の分布が少ない．これよりglobal average poolingを使った方が損失関数の凹凸が少なくなることが分かる．よって，クラス分類タスクにおいて，vision transformerの損失関数の凹凸を緩和するにはglobal average poolingを使った方が良いことが分かる．

上図はmulti-head self attentionを行う領域の広さを変化させたときの結果である．カーネルサイズとはmulti-head self attentionを行う領域の広さのことである．
図(a)はカーネルサイズの違いによるテストデータを誤識別した割合と学習時の損失関数の値を示している．これより，カーネルサイズは大きすぎず，小さすぎないサイズが良いということが分かる．
図(b)はカーネルサイズの違いによる損失関数のヘッシアン最大固有値の分布を示している．カーネルサイズが大きすぎず，小さすぎない時に負の固有値の分布が少なくなることが分かる．
これより，広い範囲や狭い範囲に対してmulti-head self attentionを行うよりもその中間の範囲に対してmulti-head self attentionを行う方が良いことが分かる．よって，Swin Transformerのようにmulti-head self attentionを行う領域を制限すると良いことが分かる．

multi-head self attentionはローパスフィルタの役割を持っている

multi-head self attentionはローパスフィルタの役割を持っており，convolution neural networkはハイパスフィルタの役割を持っている．

上図は$ViT$における層の深さと高周波成分の相対対数振幅である．高周波成分の相対対数振幅とは，その特徴マップにどれくらい高周波成分の特徴が含まれているのか示すものである．灰色のところはmulti-head self attentionであり，白色のところはMLPである．multi-head self attentionの直後は高周波成分の特徴が小さくなり，MLPの直後は高周波成分の特徴が大きくなる．これより，multi-head self attentionは高周波成分の特徴を小さくするため，ローパスフィルタの役割を持っていると言える．

上図は$ResNet$における層の深さと高周波成分の相対対数振幅である．白色のところはconvolution neural networkであり，青色のところはプーリング層である．convolution neural networkの直後は高周波成分の特徴が大きくなっている．これより，convolution neural networkはハイパスフィルタの役割を持っていると言える．