TransGAN: Two Transformers Can Make One Strong GAN

[Yagami360]

0. 論文情報・リンク

• 論文リンク：https://arxiv.org/abs/2102.07074
• 公開日時：2021/02/14
• 組織 : MIT-IBM Watson AI Lab
• 被引用数（記事作成時点）：8 件
• 実装コード（公式：PyTorch）: https://github.com/VITA-Group/TransGAN
• Publication :

1. どんなもの？

• · NLP の分野で優れた性能を発揮している Transformer を、noize-to-image の GAN モデルに適用できる形にした TransGAN を提案。既存の CNN ベースのアーキテクチャの GAN モデルと比較して競争力のある品質を実現
  

2. 先行研究と比べてどこがすごいの？

• self-attention の構造をもつ Transformer は、自然言語処理（NLP）の分野で目覚ましい成果を収めている。
  一方で、画像分類の分野でも、Vision Tranfomer のような Transfomer の構造をもったモデルが、従来の CNN ベースのアーキテクチャと同等かそれ以上の性能を発揮し始めている。
  本モデル（TransGAN）は、画像分類より困難な生成系の分野でも Transformer を適用する試みになっており、従来の CNN で構成されていた noize-to-image の GAN アーキテクチャを、全て Transformer に置き換えた初のモデルになっている。
  この際、Vision Tranfomer のようなアーキテクチャを生成系にそのまま適用しても、画像パッチのシーケンス長が長くなりすぎてメモリや計算コストの問題が発生するので、PGGAN のように段階的に解像度を高くしながら画像生成を行う工夫や Co-Training の工夫、self-attention の局所化などの工夫を取り入れて、既存の CNN ベースのアーキテクチャのGAN モデルと比較して競争力のある品質を実現している。

3. 技術や手法の"キモ"はどこにある？

• アーキテクチャの全体像
  

• 生成器
  
  TransGAN では、上図の Transfomer Encoder の部分のように、Transformer を画像系の分野に適用した Vision Tranfomer（ViT）を基本的な構造として取り入れている。
  但し、Vision Tranfomer のようなアーキテクチャを生成系にそのまま適用すると、たとえ低解像度の画像（32×32など）であっても、入力する画像パッチは長いシーケンス（1024）となり、メモリ消費量が爆発的に生じてしまう。また、Transfomer encoder 内での attention 部分での計算コストも爆発的に増加してしまう。
  そのため TransGAN では、PGGAN のように段階的に解像度を高くしながら画像生成を行う工夫を取り入れている。

  具体的には、以下の処理を行う

  1. 通常の noize-to-image の GAN モデルと同じ用に、ランダムな入力ノイズを生成する。
  2. 入力ノイズを MLP に入力し、8x8xC のベクトル長をもつベクトルに変換し、64 個の長さ C のベクトルに flatten 化する。
  3. 64 個の長さ C のベクトルを位置埋め込みを加算した上で、ViT に入力する。
  4. ViT からの出力特徴マップを UpSampling 層（pixelshuffle +reshape）に入力し、解像度を２倍（16x16xC/4）にする。
  5. 同様の処理を 32x32 の解像度になるまで繰り返し、最後に全結合層（Liner Unflatten）に入力して reshape することで最終的な 32x32x3 の画像を生成する。

  32x32 の解像度までの画像生成を行う場合の処理を記載

  本論文では、32x32 or 64x64 の解像度までの画像生成しか行っておらず、64x64x3 でも品質があまり良くないことから、より高解像度（1024など）での画像生成の品質は、かなり悪いと思われる。

  生成器では、ViT のように入力画像から分解した画像パッチを入力するのではなくて、ランダムノイズを MLP に通した後で reshape したベクトルを入力している点に注意

• 識別器
  
  GAN の識別器は、入力された画像に対して｛本物 or 偽物｝のラベル情報を出力すればいいので、ViT と同じく入力画像をより細かい画像パッチに分解し、画像パッチを Transfomer に入力するようにする。
  また画像パッチを flatten 化した後に、flatten 化された画像パッチに対して cls トークン（上図では 0 の部分）を concat で付与している。
  そして、cls トークンに対応した出力を全結合層（Head）に入力することで、識別器の出力（本物/偽物ラベル）が得られるようにしている。

  cls トークンは、ランダムな初期値から始まる学習可能なパラメーターになっている。

• データオーギュメントの効果
  

  Tranfomer や ViT は、CNN のような位置不変性の性質を持たないために、学習用データセットの数が非常に必要なモデルになっている。
  TransGAN に対してもこのことは同様であるので、学習用データセットのデータオーギュメント（DA）は非常に有効になると考えられる。
  上図は、実際に CNN ベースの既存の GAN モデルと TransGAN に対して、DAありなし（DA手法は DiffAug に基づく）での品質スコアを比較した図であり、TransGAN では CNN ベースの GAN モデルよりも更に DA が有効になっていることがわかる。

• Co-Training（共同学習）
  
  Co-Training（共同学習）は、本来のタスクとは別のタスクでもモデルを自己教師あり学習で学習させることにより、学習を安定化させる方法である。
  TransGAN では、上図のような超解像度化タスクでもモデルを学習させることにより、GANでの学習を安定化させている。
  ここで、高解像度画像は元の入力画像から得られ、低解像度画像は元の入力画像をリサイズしたものから得られるので自己教師あり学習で学習できる。また超解像度化タスクでの損失関数は、生成高解像度画像と正解高解像度画像との間の MSE となり、その損失関数を生成器の損失関数に線形結合で追加すればよい。

  

  上図は、実際に Co-Training で学習した場合の効果を示した図である。
  Co-Training で学習した場合（MT-CT）、生成画像の品質スコアが向上しているがわかる。

• Locality-Aware Initialization for Self-Attention
  

  CNN では、上図のように画像の隣接ピクセルの情報しか伝搬できないのに対して、self-attention 構造をもつ Transfomer では画像全体のピクセルからの情報を伝搬できるという特徴があり、これが Transfomer を画像タスクに適用する場合の優位性になっている。

  

  一方で、Transfomer のみで構造された ViT 内部での特徴マップは、上図のように、CNN 内部の特徴マップと同じような縞模様の特徴を重点的に学習している性質が確認されている。
  このことからも CNN のもつ局所性というのは、画像タスクに対しては重要な性質になっていると考えられる。

  

  CNN を完全排除して Transfomer のみで構造した TransGAN にこのような局所的を持たせるために、本手法では上図のように、学習の初期段階ではマスク（上図灰色）されていない領域のみで self-attention するようにして、学習の最終段階に近づくにつれて、マスク領域を徐々に減らしていくような Self-Attention 構造を導入している。

  

  上図は、実際にこの処理（Locality-Aware Initialization for Self-Attention）の効果を示した図である。
  この処理を有効化した場合（LOCAL INIT）に、生成画像の品質スコアが向上しているがわかる。

4. どうやって有効だと検証した？

• 学習用データセット

  • CIFAR-10 : 32x32 の物体or動物画像
  • CelebA : 64x64 の人物顔画像

• 実験設定
  xxx

• 既存の noize-to-image GAN モデルとの定量的品質比較
  

  CIFAR-10（32x32）で学習した場合においては、TransGAN での品質スコアが WGAN-GP よりも上回っており、それなりに競争力のあるモデルになっている。

  そのそも 32x32 の解像度しかない CIFAR-10 ではそんなに IS, FID スコアの差異がでない。より大きな解像度でやった場合は TransGAN のスコアがかなり悪くなりそう。

• 本モデルでの生成画像の品質
  

  CIFAR-10 や CelebA で学習した場合に、TransGAN ではそれなりの品質の生成画像が生成できている。

  より大きな解像度（1024）での品質が気になる。64x64 の CelebA でも劣化していることを見ると、高解像度での品質はかなり悪そう。

5. 議論はあるか？

• 32x32x3 or 64x64x3 での解像度までの画像生成しかやっておらず、64x64x3 でも品質があまり良くないことから、より高解像度（1024など）での画像生成の品質は、かなり悪いと思われる。
• 現時点では、論文で主張されているような TransGAN が既存の CNN ベースの GAN と比べて競争力があるとまでは言いがたい。この論文の成果としては、GAN において CNN を全て Transfomer に置き換えた初のモデルであって、今後の Transfomer ベースの GAN 研究の指針になりえるという点。
• UpSampling 層での pixelshuffle は、PyTorch では nn.PixelShuffle(scale_factor=2) で簡単に実装できる。pixelshuffle 自体はかなり前に提案されているアップサンプリング手法だが、nn.ConvTranspose2d や nn.Upsample での UpSampling 層よりも生成画像の品質が向上するか？
• GAN の学習を安定化させるテクニックとして、Co-Training で超解像化タスクを自己教師あり学習で学習する方法は初めて知った。他のGANモデルでも有効そう

6. 次に読むべき論文はあるか？

• xxx

7. 参考文献

• https://qiita.com/omiita/items/35ea4543bc3e52989a69
• [Transformer] Attention Is All You Need
• [Vision Transformer] An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale