【論文メモ】 A U-Net Based Discriminator for Generative Adversarial Networks

A U-Net Based Discriminator for Generative Adversarial Networks を読んだ。

A U-Net Based Discriminator for Generative Adversarial Networks

（以下、言及のない図は↑から引用しています）

CVPR 2020 に accept されていて、arXiv への公開も2020年2月28日と最近。結果がすごくて、3つのデータセットで SoTA を達成している。

提案されているのは以下の２つ。

・U-Net base の Discriminator
・CutMix data augmentation による consistency regularization

Discriminator で global な特徴を使用するために U-Net を採用し、CutMix data augmentation で consistency regularization を行うという趣旨のよう。
Multiscale の情報を利用した GAN 自体はこれまでもあったけど、この論文はとても結果がいいので読んでみました。

U-Net 自体は MICCAI (Medical Image Computing and Computer-Assisted Intervention) 2015 という医療系の会議で発表された U-Net: Convolutional Networks for Biomedical Image Segmentation という論文（12000↑ citations！）で初めて提案された手法で、以下のような特徴を持ちます。

・左右対称な Encoder-Decoder 構造
・Encoder の特徴マップを対応する Decoder の特徴マップに concatenate する

アーキテクチャ図は↓（ U-Net: Convolutional Networks for Biomedical Image Segmentation より引用）。

要は異なる解像度での特徴を skip で連結している。こういう特徴があるので、global, local 両方の data representation を獲得することができるのでは、というのは昔から言われていて、Segmentation で付近の情報だけでなく大域的な情報を使う必要があったりする問題設定などに有効だと知られていた（U-Net の元論文のアプリケーションも電子顕微鏡画像からの細胞膜セグメンテーションだった）。

Generator を U-Net にした GAN はだいぶ前からあり、有名なものに PFN の PaintsChainer がある。ちなみに、2017年の僕の修士論文も Generator を U-Net にした cGAN で電子顕微鏡画像を生成していたりする。
こんな感じで、2017年頃当時は cGAN も U-Net も新しく色々試されていたものの、2020年になって U-Net を Discriminator として採用してここまで圧倒的な結果が出るというのは意外だった。U-Net のポテンシャルの高さを表していると思う。

D^U_enc と D^U_dec があるのが特徴。それぞれ以下のような loss になる。

encoder 側については一般的な GAN の Discriminator の loss と同じ。
decoder 側の loss は面白くて、pixel index i, j で平均を取っている。この pixel 毎の出力は skip connection で持ってきた global な特徴と encoder の出力の upsampling で持ってきた local な特徴を含んでいる。

続いて、CutMix data augmentation と Consistency Regularization について。
CutMix 自体は CutMix: Regularization strategy to train strong classifiers with localizable features, ICCV 2019 で提案された手法で、異なるクラスの画像を切ったり（cut）混ぜ合わせたり（mix）する data augmentation です。

real image と fake image が混合されるように、real 画像に対して ratio に従ったマスクをしてマスクをした箇所に fake image を貼る・・・ということをしています。画像分類でやる random erasing augmentation にちょっと似てる。
実装では、ミニバッチを作るときに一定確率に従って CutMix 画像を加えているっぽい（この確率は epoch ごとに上がっていく）。

こうして作った画像（ mix(x, G(z), M) と表す）を使って、↓のようなペナルティ項を作る。

CutMix 画像を使った Discriminator の出力と decoder を通した画像の CutMix が近くなるようにすることで、意味的・構造的な変化に強くなるようにしている。

結果は、顔画像の FFHQ、動物画像の COCO-Animals、有名人の顔画像の CelebA でそれぞれ FID と IS で比較、BigGAN を上回り SoTA。

感想

GAN ではこれまで local, global 両方で coherent な画像を生成するのは難しかったように思うけど、Discriminator への U-Net の採用 + 正則化というシンプルなアイディアで解決していてすごい。実装簡単そう