Vision Transformer (ViT)

Vision Transformersは、画像をパッチに分割し、パッチの1次元表現の並びを、標準的なTransformerが単語の並びを扱うのと同じように扱うTransformerです。

Vision Transformer

1.bp.blogspot.com

Vision Transformerについては、Google AI blogに詳しい説明があります。

ViTGANのアーキテクチャ

ジェネレータは以下のアーキテクチャに従っています：

ViTGAN ジェネレータのアーキテクチャ

drive.google.com

ジェネレータは、標準的なVision Transformerとは以下の点で異なります：

• Transformerへの入力は、位置のエンベディングのみで構成されています。
• LayerNormの代わりにSelf-Modulated Layer Norm (SLN)を使用する。$z$で定義されるシードがネットワークに影響を与えるのはここだけである。
• 分類ヘッドの代わりに、各パッチのエンベッディングが$\textup{SIREN}$を通し、各パッチの画像が生成される。パッチをつなぎ合わせて、画像を生成します。

---

ディスクリミネータは以下のアーキテクチャに従っています：

ViTGAN ディスクリミネータのアーキテクチャ

drive.google.com

ViTGAN ディスクリミネータは、標準的なVision Transformerネットワークに、以下のような変更を加えたものです。

• DiffAugment。
• 重複する画像パッチ。
• ベクトル化されたL2距離によるself-attention。
• 改良型スペクトル正規化(ISN)。
• バランスド・コンシステンシー正則化（bCR）。

ディスクリミネータにおけるリプシッツ性の確保

リプシッツ連続性とは、関数で、その勾配に限界がある事：

$$\bigg | \frac{df(x)}{dx} \bigg | \leq C, \qquad C \geq 0$$

勾配に制限がないと、学習が非常に不安定になります。1回のトレーニングステップでモデルの重みが無限に変化する可能性がある。そのため、勾配の境界を確保するといいです。

ディスクリミネータにおける、ベクトル化されたL2距離によるself-attention

ドット積を用いた標準的なself-attentionは、最近、境界のない勾配を持つことが示されました [paper]。

ViTGANのディスクリミネータは、この最近提案されたユークリッド距離を用いたself-attentionを採用しています、（$d(\cdot,\cdot)$）。

一般的なself-attention

$$Attention_h(X) = softmax \bigg ( \frac{QK^T}{\sqrt{d_h}} V \bigg )$$

リプシッツ性のself-attention

$$Attention_h(X) = softmax \bigg ( \frac{d(Q,K)}{\sqrt{d_h}} V \bigg )$$

改良型スペクトル正規化

以下のようなスペクトル正規化：

$$\bar{W}_{SN}(W):=W/\sigma(W)$$

を使用すると、ディスクリミネータのリプシッツ連続性をさらに強化します。リプシッツ定数が1になることが保証されます。つまり、勾配は常に$1$または$-1$になります。

しかし、ViTGANの作者は次の事実に気づきました。Transformerがリプシッツ定数に非常に敏感であり、定数が小さいと学習速度が遅くなり、情報が崩壊してしまう事もあります。

その理由で、改良型スペクトル正規化 (ISN)が導入されています：

$$\bar{W}_{ISN}(W):=\sigma(W_{init})\cdot W/\sigma(W)$$

これにより、異なる層の勾配は、モデルの初期化時の値と同じ、異なるリプシッツ定数を持つことになります。

重複する画像パッチ

ディスクリミネータは過剰適合しやすい。

一般的ななビジョンTransformerでは、あらかじめ定義されたサイズのグリッドに従って、画像が分割されます。

画像	パッチ (8x8)
トカゲの画像	トカゲの画像のパッチ

しかし、グリッドサイズの選択は非常に重要で、全く異なる結果をもたらします。

問題は、ディスクリミネータが局所的な手がかりを記憶してしまい、意味のある情報がジェネレーターに伝わらず、学習が止まってしまいます。

ViTGANのディスクリミネータは、重複する画像パッチを導入しています:

画像	重複する画像パッチ(16x16)
トカゲの画像	トカゲの画像の重複するパッチ

一般的なパッチの代わりに重複するパッチを使用することで、以下の2つの点でTransformerを改善することができます：

• パッチグリッドの大きさの影響を受けにくい。
• 隣接するパッチ間でピクセルが共有されるため、位置関係の認識が向上します。

DiffAugment

移動、ランダムクロッピング、カラーシフトなどの伝統的な画像補強技術は、実際のデータセットにのみ適用されます。これらの技術は、ディスクリミネータをよりロバストにし、入力のわずかな変化に影響されないようにするために使用され、結果的に過剰適合を緩和するのに役立ちます。

微分可能な増強(DiffAugment) [Paper]は、情報を逆伝播させることができるように画像を増強する方法です。DiffAugmentにより、「本物」のデータセットだけでなく、生成された「偽物」の画像にも増強を適用することができます。

通常の画像増強の代わりにDiffAugmentを使用すると、非常に大きな差で結果が改善されることが証明されています。

SLN ・ Self-modulated LayerNorm

ViTGANのジェネレーターは、下記のように計算される、Self-modulated LayerNormを使用しています。

$$SLN(h_{\ell},w)=\gamma_{\ell}(w)\odot\frac{h_{\ell}-\mu}{\sigma}+\beta_{\ell}(w)$$

$\gamma_\ell(w)$と$\beta_\ell(w)$は1層のネットワークで、潜在値 $z$ に基づいて、層のバイアスと偏差を計算します。

$\textup{SIREN}$

$\textup{SIREN}$は、$sine$の活性化関数を使ったネットワークで、自然な連続信号の暗黙的な表現を学習するのに使われます。詳しい説明は、Sinusoidal Representation Networks SIREN - Fusic Tech Blogをご覧ください。

基本的に、$\textup{SIREN}$は単一の画像を表現するように学習されますが、ViTGANの作者は、フーリエ位置埋め込みにパッチ埋め込みを組み込むことに成功しました。それで、$textup{SIREN}$は様々な画像を表現できるようになっています。

$E_{fou}$フーリエ位置エンベッディング

フーリエ位置エンベッディングは、Transformerの位置エンベッディングと同様に、$sin$の波に従うものです。以下の式で算出されます。

$$E_{fou}(\mathbf v)= \sin(\mathbf W \mathbf v)$$

$\mathbf v$は画素の$x$と$y$軸の位置のベクトルであり、$-1$と$1$の範囲に入るように正規化をします。$\mathbf W$は学習する重みです。2次元の$\mathbf v$を適切な大きさのベクトルの射影の重みになります。

ウエイトモデュレーション

ViTGANのジェネレーターはウエイトモデュレーションを使用しています。ウエイトモデュレーションとは、2つの入力チャンネルを合わせるための非常に効率的な道具です。StyleGAN2やCIPSと同じ手法で位置エンベディングとパッチのエンベディングを合わせます：

$$w^{'}_{ijk}=s_i\cdot w_{ijk}$$

ウエイトモデュレーションを使用した場合の結果は、単純な足し算や結合の手法よりいいです。スキップコネクションを使用したモデルよりもいい結果が得られています。

バランスド・コンシステンシー正則化(bCR)

ViTGANではbCR[Improved Consistency Regularization for GANs]が使用されています。

bCRとは、与えられた画像とその増強された画像に対して、同じ予測値を出力するようにディスクリミネータを学習する事です。

バランスド・コンシステンシー正則化

drive.google.com

画像を少し移動させたり、色を少し変えたり、その他どんな小さな補強を加えても、ディスクリミネータの出力は、補強に関わらず、同じであってほしいと思います。そうでなければ、入力に小さな変化があっても、ディスクリミネータの出力は大きく異なってしまいます。

bCRを使用することで、ディスクリミネータとジェネレーターの両方の学習を安定させることができる。

まとめ

私の意見では、ViTGANは非常によく考えられているモデルである。

作者はTransformerの学習における弱点に多くの焦点を当てています。特に、リプシッツ連続性を保証することで、結果の質が大幅に向上することが証明されています。

結果の改善に最も貢献した要因は、DiffAugmentの使用であると思われます。

様々なディスクリミネータの比較

Kwonjoon Leeらは、Transformer、特にVision Transformerを学習する際の弱点を良く研究したと思います。 sformer.

Vision Transformerは2020年末から存在していましたが、Kwonjoon LeeらはGANに適用した最初のチームの1つです。Vision Transformer GANに関する最初の出版物で、このような素晴らしい結果を達成したことは信じられません。さらに、彼らは斬新なアプローチで$\textup{SIREN}$の活用に成功しています。