TransGAN: Two Transformers Can Make One Strong GAN

[IsHYuhi]

INFO

author

Yifan Jiang¹ Shiyu Chang² Zhangyang Wang¹

affiliation

¹Department of Electronic and Computer Engineering, University of Texas at Austin, Texas, USA
²MIT-IBM Watson AI Lab, Massachusetts, USA

conference or year

2021

link

arXiv
実装

概要

近年のComputerVision分野においてもtransformerが強力な'universal'モデルとなる可能性が示唆されている. しかし敵対的生成ネットワークのような難しいvisionタスクを扱うことができるのか, この研究では完全にCNNを用いないTransformerベースのGANを構築するpilot studyを行った. TransGANはData Augmentation, multi-task co-training strategy, 画像の近傍平滑性を強調するloccally initialized self-attentionにおいて効果が顕著であることを実証し, CNNベースのbackboneに基づくSoTAモデルに匹敵する性能を達成.

GANは(Goodfellow+, 2014)は全結合層のみで構築されており, 小さな画像しか生成できなかった. その後DCGAN(Radford+, 2015)がCNNを用いてより高画質で深い生成モデルの安定した訓練を可能にした. そこからほとんどのGANは CNN-basedなGeneratorとDiscriminatorに依存していおり, 確かに畳み込みは多様で魅力的なvisual結果に貢献している.

基本的に畳み込みは局所的な受容野を持つので, CNNは十分な数のレイヤーを通過しないとlong-rangeな依存性を処理できない. それだけだと最適化が難しくなるだけではなく, 特徴の解像度や細かいdetailが失われる可能性がある. 従来手法でSelf-Attentionやnon-local operationなど導入することによる利点から実証されているようにCNNはglobalな統計を捉えることに本質的には適していない.

Computer VisionタスクにおいてTransfromer の利点は２つ挙げられる.

• 強力な表現力を持ち, 人間が定義したinducive biasがない. CNNは特徴の局所性に強いバイアスを持ち, フィルタの重みを全ての場所で共有することで空間的な不変性を持つ.
• transformerのアーキテクチャは一般的で, 概念もシンプルで, タスクやドメインを超えた強力な'universal'モデルになる可能性を秘めている. (CNN-basedでみられるad-hoc building blocks (Res-Netなどを参照)を取り除くことができる.)

課題

image patcheのシーケンスに直接適用されるpure transformer architectureは, 画像分類タスクで非常に良い性能を発揮することができるが, 構造, 色, テクスチャの空間的な一貫性に対する高い要求がある画像生成の場合, 同じ方法が有効であるかは不明である. 画像を出力する既存のtransformerモデルは一様にCNN-basedなエンコーダーやCNNでの特徴抽出を用いている. しかし, よく設計されたCNNベースのアーキテクチャだとしても, GANの訓練は不安定で, モード崩壊を起こしやすい. 他にもtedious, heavy, data-hungryであることが知られている. これらの２つをminglingすることで訓練時の課題が増幅されることは疑いようがない.

単純に考えれば, そのままピクセルの入力から複数のtransformerブロックを積み重ねることができるが, メモリや計算量が膨大になる. その代わりに, 各ステージでembedding次元をを各ステージで減らしながら, 徐々に特徴マップの解像度を上げていく. このTransGANはself-Attentionのグローバルな需要やの恩恵を受けているが, 実際は生成能力の低下や視覚的な滑らかさの崩壊が起きる. CNN-based GANとの性能差を縮めるために, Data Augmentation, multi-task co-training with self-supervised auxiliary loss, localized initialization for self-attentionから顕著な恩恵を受けることを実証する.

既存研究

Taming TransformerではCNNGANを用いて, 文脈に富んだvisual partsのコードブックを学習し, その合成はその後, autoregressive transformer architecturでモデル化される. しかし, しかし全体的なCNNアーキテクチャはそのままで, codebookや量子化なども,　モデルの汎用性を制限している. To our best knowlede, CNNを完全に除去した既存研究はない.

提案手法

encoderとしてViTを選び, 変更は最小に抑える(もしくは, 最小限の変更を加える？とあるが相違点はどこ?). encoder はmulti-head self-attention moduleとfeed-forward MLP with GELU non-linearityからなる.
低解像度の画像（例えば32×32）であっても, 長いシーケンス（1024）になり, 爆発的なself-attentionのコスト（シーケンス長の2次関数）が発生する. このコストを回避するために, 複数の段階で解像度を反復的にup scaleする.(入力シーケンスを徐々に増やし、埋め込み次元を減らす). UpScaling層ではPixel Shuffleを行っている.

(32x32x3)の画像を生成する場合:

(Noize) --MLP-> (64x64) --encxM-> (64x64) --reshape-> (8x8x64) --Pixel Shuffle-> (16x16x16) --reshape-> (256x16) -
-encxL-> (256x16) --reshape-> (16x16x16) --Pixel Shuffle -> (32x32x4) --reshape-> (1024x4) --encxL-> (1024x4) -
-Linear Unflatten-> (32 x 32 x 3)

DiscriminatorはViTと同じ構造で(8x8)のパッチに分割Layerは7層. [0]が[cls]トークンに該当. ((N+1)xC)

Multi-task co-training with self-supervised auxiliary task

BERTが自己教師あり学習によって性能向上が得られていることにインスパイアされている.
Stage 2に低解像度の画像を入力し, 高解像度の画像を得るタスクを同時に学習させる. Loss FunctionはMSE loss.

Locality-Aware Initialization for Self-Attention

ViTが画像からCNNと同じような特徴量を学習することが報告されている.
ここではSelf-Attentionに局所性を意識させた初期化を導入する. 実際にはマスクを訓練中に徐々に減らし, みれる範囲を広げ, 最終的には完全なSelf-Attentionになる.これはSelf-Attentionは初期の訓練段階で役立つが, 後半の訓練段階では最終的に達成可能なパフォーマンスを損なう可能性があるという観察に基づいて作成されている.

検証

Transformer vs CNN

TransformerベースのGeneratorとDiscriminatorの性能を比較するため, CNN-basedでSoTAなAutoGANと組み合わせ検証.
AutoGAN: Neural Architecture Search(NAS)で作られたGAN.
NAS: AutoMLとかで使われてるやつ.

• (TG+AD)TransformerベースのGeneratorは, AutoGANのDiscriminatorでトレーニングした場合, その性能はすでにオリジナルのAutoGANと同等になるという強い能力を持ってる. これは, GeneratorにTransformersを入れるとうまくいくというTaming Transformer(Esser et al., 2020)とも一致している.

• (AG+TD)しかし, TransformerベースのDiscriminatorは競争相手としては劣っており, AutoGANのGeneratorを良い生成結果に導くことはできないようである.

• (TG+TD)AutoGANのGeneratorをTransformerベースのGeneratorに置き換えると, GとDの構造がより対称的になったことが功を奏したのか, 結果が若干改善. しかし, 畳み込みDを使用した場合には, まだ大きな遅れをとっている.

Data Augmentation

MT-CT, LOCAL INIT.

Scaling up

xM, xLの部分に該当 {xM, xL, xL}
スケールアップさせることで性能が向上. Discriminatorのスケールアップによるゲインはあまりなかった.

定性評価

定量評価

新規性

Model Architecture:

純粋にTransformerのみを使い, 畳み込みを行わないGANを構築した. メモリのオーベーヘッドを避けるために, メモ入りに優しいGenerator と Discriminatorを作成し, より大きなモデルに効果的にスケールアップが可能.

Training Technique:

TransGANをうまく訓練するためにData Augmentation, multi-task co-training with self-supervised auxiliary loss, localized initialization for self-attentionを研究.

Performance:

TransGANはCNNベースのものと匹敵する性能を達成.

議論，展望

We point out a few specific items that call for continuing efforts:

• More sophisticated tokenizing for both G and D, e.g. using some semantic grouping (Wu et al., 2020).
• Pre-training transformers using pretext tasks (Dai et al., 2020), which may improve over our current MT-CT.
• Stronger attention forms, e.g., (Zhu et al., 2020).
• More efficient self-attention forms (Wang et al., 2020; Choromanski et al., 2020), which not only help improve the model efficiency, but also save memory costs and hence help higher-resolution generation.
• Conditional image generation (Lin et al., 2019).

Comment

date

May 20th, 2021