[論文メモ] UVCGAN: UNET VISION TRANSFORMER CYCLE-CONSISTENT GAN FOR UNPAIRED IMAGE-TO-IMAGE TRANSLATION

Vision Transformer(ViT)を使ったunpaired image-to-image translation。

unpaired image-to-image translationでは既存手法としてCycleGAN、ACL-GAN、Council-GANなどがある(U-GAT-ITも)。
CycleGANはone-to-oneだが、ACL-GANとConucil-GANはone-to-manyで多様性に重きをおいている。

現代のニューラルネットアーキテクチャにCycle consistency lossを使ってone-to-one translationを改善できるのでは？というお気持ち

手法

提案するUNet ViT Cycle-consistent GAN(UVCGAN)のアーキテクチャ、学習フローは以下の図1。
f:id:Ninhydrin:20220315091840p:plain

2つの画像ドメイン $A, B$ を相互に変換するGeneratorをそれぞれ $\mathcal{G}_{A \rightarrow B}, \mathcal{G}_{B \rightarrow A}$ とする。
そして、どちらのドメイン由来かを識別するDiscriminatorを $\mathcal{D}_A, \mathcal{D}_B$ とする。

アーキテクチャはUNetをベースとし、ボトルネック部分にピクセルレベルのViTを導入。pre-process部分はconv + leakyReLU、post-process部分は1x1 conv + sigmoid。

ViT部分ではピクセルレベルのパッチ化をしてpositional encodingをconcatしてる。足さなかった理由があるのだろうか。

事前学習とloss

Bidirectional Encoder Representations from Transformers (BERT)-like self-supervised pre-trainingを行った。
BERT likeな事前学習はランダムに初期化された大規模ネットワークの事前学習に効果的で、マスクされた画像から元画像を予測するタスクを行う(ピクセルレベルの $l_1$ loss)。

事前学習に本タスクの学習を行うが最終的なlossは以下。
f:id:Ninhydrin:20220315094533p:plain

$\mathcal{L}_{GAN}$ はadversarial loss、 $\mathcal{L}_{idt}$ はidentity loss、 $\mathcal{L}_{cyc}$ はcycle-consistency loss。Discriminator用lossの右辺第二項目はGradient penalty(GP)。 $\lambda_{XX}$ はそれぞれハイパラ。

実験・結果

データセットはselfie⇔anime、male⇔femail、メガネの有無。
BERT likeな事前学習では256x256の画像を32x32ピクセルのパッチにして40%をマスクする。
f:id:Ninhydrin:20220315095326p:plain
f:id:Ninhydrin:20220315095420p:plain

Ablation studies。
BERT likeな事前学習で用いるデータセットについて

データ数は少ないが本タスクに近いデータセット
データ数は多いが本タスクより多様なデータセット

のどちらがいいかを調べるとともに、identity lossとGPの有無についても調べた。データセットはmale-to-femaleとselfie-to-animeの2つ。
Noneは事前学習なし。
f:id:Ninhydrin:20220315095450p:plain

identity lossとGPは組み合わせると改善
GPなしidentity lossありはなんとも
事前学習はGPとidentity lossと組み合わせると向上につながる

事前学習+identity lossはそうでも無いが、+GPは大きな改善もたらす。
これはGeneratorは事前学習したパラメータに対してDiscriminatorは初期値のままなので、意味不明な情報をGeneratorに流して事前学習済みパラメータを破壊するからではと予測。
もしそうならDiscriminatorも事前学習すればもっと改善するかも...とのこと。

所感

久々のunpaired I2I translation。U-GAT-IT以降あまり追ってなかった。
この間のdenoisingでもそうだったがUNetのボトルネックにAttentionモジュールを組み込むアーキテクチャが流行っている？パフォーマンスが良い？
ninhydrin.hatenablog.com

機会があればUNet+ViTをいろんなタスクに試して見たいところ。