[論文メモ] StyleGAN-XL: Scaling StyleGAN to Large Diverse Datasets

StyleGAN3でImageNetのような多様な画像を生成できるようにした。

SIGGRAPH 2022

StyleGANはハイクオリティかつある程度の制御可能なGANだが、制御可能にできる分、構造が一定でない分散の大きいデータセットの学習には適さない。提案するStyleGAN-XLはこれを解決してImageNetのような分散の大きなデータの生成を可能にした。

既存手法ではImageNetを学習できるものもあるが、少なくともStyleGAN3をナイーブに学習しても厳しい(図1参照)。
特に高解像度の生成(1024x1024)になると安定しない。
目的はStyleGAN3でImageNetを生成できるようにすること。「生成できる」とはISやFIDといった既存の評価指標で定義。

手法

既存のStyleGAN3のデフォルト設定をConfig-AとしてConfig-Fに至るまで変更を加えていく。

正則化

StyleGAN3ではstyle mixingやpath lengths regularizationは廃止された。
特にpath lengths regularizationは分散の大きなデータセットには不向きと報告されている。
しかし、モデルの学習が十分に進んだ後（200k枚くらい学習したら）に適用するとこれを回避できることがわかった。

Discriminatorにはspectral normalization(gradient penaltyなし)と、始めの200k枚の入力には $\sigma=2$ のガウシアンフィルタを適用する。
Discriminator blurringはStyleGAN3-Rの学習で利用され、Discriminatorが高周波に注目するのを防ぐ。

低次元の潜在空間

ProjectedGANの適用を考える。ProjectedGANについては以前のメモを参照。
ninhydrin.hatenablog.com

ProjectedGANはStyleGANよりもFastGANの構造に有効で、その大きな違いは潜在空間の大きさであると考える。
潜在空間 $z$ はFastGANが $\mathbb{R}^{100}$ 、BigGANが $\mathbb{R}^{128}$ 、StyleGANが $\mathbb{R}^{512}$ と大きすぎる。
近年の研究では自然画像は実は低次元に内在しているという主張があり、ImageNetは40前後と推測されている。
512次元でProjectedGANを適用するとその高速化の恩恵を受けられなかった。
そこで $z$ を512から64に変更すると、学習が安定し高速の恩恵も得られた。
ベースモデル(Config-A)に対して上記の正則化とProjectedGANを適用したのがConfig-Bとなる。

事前学習したクラス埋め込み

Config-Bだと各クラスで同じような画像が生成される(recallが0.004と低い)。
これはProjectedGANの影響と考え、クラス埋め込みを事前学習した。
Efficientnet-lite0の最小の解像度の特徴をImageNetのクラス毎に平均し、それをconditionとして利用する。
これによりrecallが0.15に上昇。この設定をConfig-Cとする。

Progressive Growingの再導入

Progressive Growingはネットワークを徐々に大きくして高解像度の生成を可能にした手法だがアーティファクトが出る等の問題でStyleGAN2で消えた。StyleGAN3ではアーティファクトの原因はエイリアシングであるとわかり、アーキテクチャでそれを回避した。
そこでProgressive Growingを再度考える必要が出てきた。
開始の解像度を16x16として11層使って開始、最終的な1024x1024の解像度のときには39層に成長。
成長時、低解像度の層はmode collapseを回避するために固定される。
これがConfig-D。

Exploiting Multiple Feature Networks

ProjectedGANのablationではFeature Networkとして色々試しているが複数の組み合わせはない。
EfficientNet-lite0を基本として、追加の2つめのFeature Networkのタスク(クラス分類 or 自己教師あり学習)、アーキテクチャ(CNN or ViT)が与える影響を調査。
追加のCNNはタスクが教師あり or 自己教師ありのどちらでもFIDをわずかに下げるにとどまったが、ViTとの組み合わせは大きくパフォーマンスを向上。
これは「教師あり学習と自己教師あり学習で得られる特徴表現は類似している」「CNNとViTは異なる表現を学習している」という既存研究の結果を裏付けている。

これより多くFeature Networkを追加しても大きな改善は得られなかったのでEfficientNetとDeiT-MをFeature NetworkとしたものをConfig-Eとする。

なお特徴表現の参考文献は下記。

教師あり学習と自己教師あり学習の特徴表現について
arxiv.org

CNNとViTの特徴表現について
arxiv.org

Classifier Guidance for GANs

diffusion modelにクラス分類器をガイドとして導入した手法がある。
生成画像 $x$ を学習済みクラス分類器 $CLS$ に通してクラスの予測 $c_i$ を得る。そしてcross-entropy loss $\mathcal{L}_{CE} = -\Sigma^Cc_i \log CLF(x_i)$ を $\lambda$ でスケールしたものをGeneratorの追加のlossとする。クラス分類器としてはDeiT-Sを採用。
これによりISが大きく向上。
この設定をConfig-Fとする。
なお、 $\lambda=8$ ぐらいがよく、解像度が大きく無いと効果がない(> $32^2$ )。

各Configでの実験結果はひょう1を参照。

実験・結果

解像度 $1024^2$ のImageNetの生成。

所感

ImageNetレベルの分散の大きいデータをある程度生成できるのはすごい。
以前DanbooruデータセットでGANを学習した記事を読んだが、そこでもPath lengths lossをやめたとあり顔等のドメインがある程度固定されたもの以外ではそれなりにアーキテクチャやlossを考える必要がある。
ProjectedGANでCNNとViTの２つを利用するというのは使えそう。CNNとViTが獲得する特徴表現はことなることを調べた論文は読んでいたが、こういう利用方法は考えなかった。なるほど。
Progressive Growingは面倒なのでやりたくない。