はじめに

TGANv2: Efficient Training of Large Models for Video Generation with Multiple Subsampling Layersを読んだのでメモ．動画生成のGANは計算コストが大きくなってしまうため解像度の低い動画の生成しかできていなかった点を解決したという論文．

assumption

この論文では計算コストの主な原因はGeneratorの後段，出力層に近い部分にあるとしている．理由としては後段では空間的な解像度が高くなるため演算回数もメモリも大きくなってしまうためとのこと．また，GANは後段のレイヤー程線形に近い単純な表現になっていると仮定していて，その性質を利用して計算コストの高くなる後段の層を効率化するという．

Abstract map and subsampling layer

Abstract mapとsubsampling layerを導入する．Notationが面倒なので代わりにGANの目的関数を以下に．

$\displaystyle \mathbb{E}_{\mathbf{x}\sim p_d}[\ln D(\mathbf{x})]+\mathbb{E}_{\mathbf{z}\sim p_z}[\ln(1-D(G(\mathbf{z})))]$

ここではgeneratorをstract block $g^A$ とrendering block $g^R$ の二つのブロックに分けて考える．Stract blockはabstract mapと呼ばれる特徴マップを計算し，rendering blockはabstract mapからデータをサンプリングする役割を持つ．簡単に言えば，generatorの出力層に近い方の畳み込みをrendering blockと呼んで残りをstract blockとしている感じ．なのでgenerator $G(\mathbf{z})$ は次のように記述可能．

$\displaystyle \mathbf{x}=G(\mathbf{z})=\left(g^R\circ g^A\right)(\mathbf{z})$

GANの計算のボトルネックはdiscriminator側にも存在し，生成されるデータの解像度が大きくなるとdiscriminatorの計算コストも上がる．そこでこの問題を解決するためにsubsampling layer $\mathcal{S}^G$ を導入する．subsampling layerはクロップ関数（画像を適当にクロップする演算）のような何かしらの関数で，abstract mapからランダムにサンプリングを行う．すると元の解像度より低い解像度のabstract mapが得られるので後段の計算コストを下げることができるというもの．このsubsampled data $\mathbf{x}'$ を生成するgeneratorを $G'$ とすれば次のように表現可能．

$\displaystyle \mathbf{x}'=G'(\mathbf{z})=\left(g^R\circ\mathcal{S}^G\circ g^A\right)(\mathbf{z})$

Abstract mapだけでなく生成されるデータの解像度も当然落ちるので目的関数も次のように修正．

$\displaystyle \mathbb{E}_{\mathbf{x}\sim p_d}[\ln D'(\mathcal{S}^D(\mathbf{x}))]+\mathbb{E}_{\mathbf{z}\sim p_z}[\ln(1-D'(G'(\mathbf{z})))]$

ここで $\mathcal{S}^D$ は $\mathcal{S}^G$ と同じようにデータをサブサンプルする関数．このように学習されたgeneratorが推論時にも同様にデータを生成するために $g^R,\mathcal{S}^G,\mathcal{S}^D$ は次の3つの条件を満たす必要があるとのこと．

$g^R$ は元のabstract mapと解像度の下がったabstract map両方を入力可能
$\mathcal{S}^G,\mathcal{S}^D$ は入力に関して微分可能
$\mathcal{S}^G,\mathcal{S}^D$ によってサンプリングされる領域は確率的に決まる

例えば， $g^R$ をfully-convolutionalなモデルにして， $\mathcal{S}^G,\mathcal{S}^D$ を固定ウィンドウ幅のクロップとすれば上の条件を満たす．

実践的な利点として， $\mathcal{S}^G,\mathcal{S}^D$ を導入すればめちゃくちゃ高次元のデータに対しても学習することが可能．もう一点の利点としてネットワーク側のパラメータも大きくすることができる．これは入力の空間的な解像度が下がればチャネル数を大きくしてもメモリ的に大丈夫ということかと．ただし，subsampling layerを一層だけしか使わなかった場合，使わないネットワークと比較して，同一学習回数においてはパフォーマンスが低下するとのこと．ただこれはabstract mapをマルチレベルにすることで回避可能らしい．

Training with multiple subsampling layers

現状，subsampling layerの導入によって計算コストと生成データの質のトレードオフが生まれた．つまり，subsampling layerによって得られたデータの解像度が小さいほど計算は効率的に行われるが，生成されるデータの質は下がる．ただし，多重解像度かもしくはmultiple frame ratesにすれば質を落とすことなく効率的にネットワークを学習可能．

まず，DCGANやSA-GANで使われるような一般的なgeneratorを考える．通常は最初に得られる特徴マップの解像度は低くチャネル数が多い．逆に出力層に近付くほど特徴マップの解像度は上がりチャネル数が少なくなる．これの意味するところは，入力に近い層ほど出力層付近に比べ学習パラメータ数が多く計算コストが小さいということ．なので入力層付近の層にsubsampling layerを入れても計算コストにおける利点が少ないことから後段の層に対して積極的にsubsampling layerを導入したい．

ここで，generatorが次のように $L$ 個のrendering blocksと $L$ 個のabstract blocksと $L-1$ 層のsubsampling layerから構成されているとする．

$\displaystyle \mathbf{x}=\left(g_L^R\circ g_L^A\circ g_{L-1}^A\circ\dots\circ g_1^A\right)(\mathbf{z})$

この演算は， $l$ 番目のgenerator（ $l$ フレーム目の生成過程）は $l-1$ 番目のgeneratorが生成したabstract mapに対しsubsampling layerを適用して得られたabstract mapを入力として受け取ることを意味する．つまり，陽に書き下せば次のようになる．

$\displaystyle G'_1=g_1^R\circ g_1^A\\ \displaystyle G'_2=g_2^R\circ g_2^A\circ\left(\mathbf{S}_1^G\circ g_1^A\right)\\ \displaystyle \vdots\\ \displaystyle G_L'=g_L^R\circ g_L^A\circ\left(\mathbf{S}^G_{L-1}\circ g_{L-1}^A\right)\circ\dots\circ\left(\mathbf{S}_1^G\circ g_1^A\right)$

この操作は計算コストの高い後段で何度もsubsamplingすることで効率的な計算を可能にしている．ちなみにFigure 2にモデルの全体像が書いてあるのでそこを見ると大体の計算の流れはわかる．

Multiple discriminators

一回のサンプリングで複数のサンプルが得られるためdiscriminatorも複数用意する必要がある．よって $l$ 番目のサンプルを受け取るdiscriminatorを $D_l'$ とすればdiscriminatorのスコアは次のように表現可能．

$\displaystyle D'(\mathbf{x}_1',\dots,\mathbf{x}_L')=\sigma\left(\sum_{l=1}^LD'_l(\mathbf{x}_l')\right)$

Adaptive batch reduction

基本的に学習はミニバッチ単位で行われるため，ここではsubsampling layerをバッチ方向にも拡張することでさらなる計算の効率化をする．ただ，Pix2PixやVid2Vvidのようなバッチサイズがめちゃくちゃ小さくてもうまくいく手法がある一方で，BigGANで言われているようにバッチサイズは大きい方がいいという議論もある．そこでこの論文では最初のサンプルの生成は多様性が必要だが，後段はそこまで多様なマップを生成する必要がない（多分abstract mapを引き継ぐ形になっているからという点が根拠かと）と仮定．その仮定のもと初期ブロックではlarge batchで学習を行い後段ではsmall batchになるようにバッチサイズのreduceを導入したとのこと．

まとめ

どう計算コストを減らしたかのみ気になったのでモデルや学習の詳細は割愛．クロップしたマップを次へ次へと渡すと生成される画像はどんどん拡大されていくような気もするがどうなんだろうか．ただ，クロップしても動くのはクロップ後の畳み込みの回数が少なくてそんな大局的な部分を見ないからというので納得はできる．Inception scoreは先行研究からそれなりのジャンプがあるスコアを叩き出しているのでかなり有効なよう．論文にも書いてあったがこの計算コストの減らし方自体は色々応用がききそう．