機械学習研究

人工知能（AI）の急速な発展に伴い、機械学習、深層学習といった言葉をよく耳にします。この「学習」とは、いったい何を指すのでしょうか。「学習」とは、望ましい結果が得られるように機械に関するネットワークやパラメータを調整する作業といえます。もう少し数理的に表現すると、正解データとネットワークからの出力データとの誤差を最小化（最適化）するようにパラメータを調整することです。多くの機械学習手法はデータに対するモデル出力の誤差（損失関数と呼ばれる）を定義し、誤差を最小化するようにパラメータの更新（学習）を行います。 このように、機械学習と最適化はとても密接な関係があります。

深層ニューラルネットワークに現れる損失最小化問題は、一般に、以下のように表現できます。 \begin{align*} \text{Minimize } f(x):= \frac{1}{n} \sum_{i=1}^n f_i (x) \text{ subject to } x\in X. \end{align*} ただし、$f_i \colon \mathbb{R}^d \to \mathbb{R}$ $(i\in \{1,2,\ldots,n\})$ は $i$ 番目の訓練データに関する損失関数、 $X$ $(\subset \mathbb{R}^d)$ は空でない閉凸集合です。

この損失最小化問題手法を解くための深層学習手法（ここではオプティマイザーと呼ぶことにします）として、確率的勾配降下法 (SGD)とそれに基づいた強力な手法がすでに提案されています。 例えば、モーメンタム法、Adaptive Gradient (AdaGrad)、 RMSProp (Root Mean Square Propagation)、 Adaptive Moment Estimation (Adam) や Adaptive Mean Square Gradient (AMSGrad) があります。

CIFAR-10データセットを用いたニューラルネットワーク ResNet-20 の訓練結果は以下の通りとなります。オプティマイザーに関する、エポック数¹⁾に対する損失関数 $f$ のグラフです。ここで適用したオプティマイザーは経験的に良いとされている学習率²⁾を利用しています。

敵対的生成ネットワーク(GAN:Generative Adversarial Networks)の目的関数は、次のように表現できます。

\begin{align*} \min_G \max_D V(D,G):= \mathbb{E}_\boldsymbol{x} [\log D(\boldsymbol{x})] + \mathbb{E}_\boldsymbol{z} [\log (1-D(G(\boldsymbol{z})))]. \end{align*}

ただし、$G:\mathbb{R}^l \to \mathbb{R}^d \ (\boldsymbol{z} \mapsto G(\boldsymbol{z}))$は生成器で、 $D:\mathbb{R}^d \to \mathbb{R} \ (\boldsymbol{x} \ $or$ \ G(\boldsymbol{z})\mapsto [0,1])$は識別器です。
さらに、$\boldsymbol{x} \in \mathbb{R}^d$ は実画像、$\boldsymbol{z} \in \mathbb{R}^l$ はノイズで、 $G(\boldsymbol{z}) \in \mathbb{R}^d$は生成画像です。 実画像の次元$d$は、データセットによって決まっています。例えば、MNISTデータセットならば$d= 28 \times 28 \times 1$となり、CIFAR10データセットならば$d=64 \times 64 \times 3$となります。生成画像の次元は、実画像の次元と等しくなるように生成器が設計されています。ノイズの次元$l$は100や128などと設定されることが多いです。

生成器はノイズ、または潜在変数と呼ばれる$\boldsymbol{z} \in \mathbb{R}^l$を入力とし、生成画像$G(\boldsymbol{z}) \in \mathbb{R}^d$を出力します。
識別器は、実画像$\boldsymbol{x} \in \mathbb{R}^d$または生成画像$G(\boldsymbol{z}) \in \mathbb{R}^d$を入力とし、入力された画像が実画像である確率$[0,1]$を出力します。

生成器は、識別器に見破られない生成画像を作ることを目指すので、$D(\boldsymbol{x})=0, D(G(\boldsymbol{z}))=1$となるように振る舞います。
識別器は、実画像と生成画像を完全に識別することを目指すので、$D(\boldsymbol{x})=1, D(G(\boldsymbol{z}))=0$となるように振る舞います。
すなわち、生成器は$L_G=V$を最小にするように行動し、識別器は$L_D=-V$を最小にするように行動します。
具体的には、生成器と識別器はAdamなどの最適化手法を使って、自身の目的関数が最小になるようなパラメータを探索します。 これを踏まえると、GANの学習は、生成器と識別器をプレイヤーとするナッシュ均衡問題の局所的ナッシュ均衡を見つけようとすることであると言えます。

このように、GANは他の深層生成モデルと比べて高品質な画像を生成することができますが、その学習は非常に不安定であることが知られています。
より早く、より安定した学習で、より高品質な画像を生成するために、たくさんのGANが提案されています。

生成器のニューラルネットワークのパラメータを$\boldsymbol{\theta}$、識別器のニューラルネットワークのパラメータを$\boldsymbol{w}$とします。それぞれの目的関数を最小にすることを目指すので、 \begin{align} \nabla_\boldsymbol{\theta} L_G(\boldsymbol{\theta}_n, \boldsymbol{w}_n)=\boldsymbol{0}, \nabla_\boldsymbol{w} L_D(\boldsymbol{\theta}_n, \boldsymbol{w}_n)=\boldsymbol{0} \end{align} を満たすような最適解$(\boldsymbol{\theta}^\star, \boldsymbol{w}^\star)$を探します。上の等式は、 \begin{align} \forall \boldsymbol{\theta} \in \mathbb{R}^\Theta, \forall \boldsymbol{w} \in \mathbb{R}^W: \langle \boldsymbol{\theta}_n - \boldsymbol{\theta}, \nabla_{\boldsymbol{\theta}} L_G (\boldsymbol{\theta}_n, \boldsymbol{w}_n) \rangle \leq 0 \text{, } \langle \boldsymbol{w}_n - \boldsymbol{w}, \nabla_{\boldsymbol{w}} L_D (\boldsymbol{\theta}_n, \boldsymbol{w}_n) \rangle \leq 0 \end{align} と等価なので、この不等式を満たすような最適解$(\boldsymbol{\theta}^\star, \boldsymbol{w}^\star)$を探しましょう。 上の不等式の左辺について、それぞれ次の不等式が成り立ちます。 \begin{align} &\frac{1}{N} \sum_{n=1}^{N} \mathbb{E} \left[ \langle \boldsymbol{\theta}_n - \boldsymbol{\theta}, \nabla_{\boldsymbol{\theta}} L_G (\boldsymbol{\theta}_n, \boldsymbol{w}_n) \rangle \right] \leq \underbrace{\frac{\Theta \mathrm{Dist}(\boldsymbol{\theta}) H^G}{2 \alpha^G \tilde{\beta_1^G}}}_{A_G} \frac{1}{N} + \underbrace{\frac{\sigma_G^2 \alpha^G}{2 \tilde{\beta_1^G} \tilde{\gamma}^{G^2} h_{0,*}^G}}_{B_G} \frac{1}{b} + \underbrace{\frac{M_G^2 \alpha^G}{2 \tilde{\beta_1^G} \tilde{\gamma}^{G^2} h_{0,*}^G} + \frac{\beta_1^G}{\tilde{\beta_1^G}} \sqrt{\Theta \mathrm{Dist}(\boldsymbol{\theta}) ( \sigma_G^2 + M_G^2 )}}_{C_G}, \\ &\frac{1}{N} \sum_{n=1}^{N} \mathbb{E} \left[ \langle \boldsymbol{w}_n - \boldsymbol{w}, \nabla_{\boldsymbol{w}} L_G (\boldsymbol{\theta}_n, \boldsymbol{w}_n) \rangle \right] \leq \underbrace{\frac{W \mathrm{Dist}(\boldsymbol{w}) H^D}{2 \alpha^D \tilde{\beta_1^D}}}_{A_D} \frac{1}{N} + \underbrace{\frac{\sigma_G^2 \alpha^D}{2 \tilde{\beta_1^D} \tilde{\gamma}^{D^2} h_{0,*}^D}}_{B_D} \frac{1}{b} + \underbrace{\frac{M_D^2 \alpha^D}{2 \tilde{\beta_1^D} \tilde{\gamma}^{D^2} h_{0,*}^D} + \frac{\beta_1^D}{\tilde{\beta_1^D}} \sqrt{W \mathrm{Dist}(\boldsymbol{w}) ( \sigma_D^2 + M_D^2 )}}_{C_D} \end{align} ただし、$\alpha^G, \alpha^D$はoptimizerの学習率で、Nはステップ数、bはバッチサイズです。その他の定義については論文を参照してください。このことから、ステップ数$N$とバッチサイズ$b$を大きく、学習率$\alpha^G, \alpha^D$は小さくすれば、それぞれの右辺は0に近くなり、局所的ナッシュ均衡を近似できることがわかります。

局所的ナッシュ均衡を近似できる理想的なステップ数$N$、バッチサイズ$b$、学習率$\alpha^G, \alpha^D$などのパラメータを設定できたとします。このとき、$\epsilon_G, \epsilon_D > 0$を十分小さい正数だとすると、 \begin{align} \frac{A_G}{N_G}+\frac{B_G}{b}+C_G=\epsilon_G^2, \text{ } \frac{A_D}{N_D}+\frac{B_D}{b}+C_D=\epsilon_D^2 \end{align} すなわち、 \begin{align} N_G(b)=\frac{A_G b}{(\epsilon_G^2 -C_G)b -B_G}, \text{ } N_D(b)=\frac{A_D b}{(\epsilon_D^2 -C_D)b -B_D} \end{align} が成り立ちます。ただし、生成器のステップ数を$N_G$、識別器のステップ数を$N_D$とし、それぞれをバッチサイズ$b$の関数だと捉えています。$N_G(b)$を微分して、その形状について考えてみましょう。 \begin{align} \frac{\mathrm{d}N_G(b)}{\mathrm{d}b} = \frac{-A_G B_G}{{(\epsilon_G^2 - C_G)b - B_G}^2} \leq 0, \text{ } \frac{\mathrm{d}^2 N_G(b)}{\mathrm{d}b^2} = \frac{2A_G B_G(\epsilon_G^2 - C_G)}{{(\epsilon_G^2 - C_G)b - B_G}^3} \geq 0 \end{align} が成り立ちます。1階微分が$0$以下で、2階微分が$0$以上なので、$N_G(b)$は$b$に対して単調減少で、かつ凸関数であることがわかります。$N_D(b)$についても同様です。
このことから、学習が収束するまでに必要なステップ数$N$を最小にするバッチサイズ$b$が存在することが分かります。

クリティカルバッチサイズとは、学習の計算量を最小にするバッチサイズのことです。 ニューラルネットワークの訓練には、クリティカルバッチサイズが存在することが先行研究によって明らかになっています。 そこで、GANの訓練にもクリティカルバッチサイズが存在するかどうかを考えてみましょう。

学習の計算量の指標である確率的勾配計算コスト（SFO計算量：stochastic first-order oracle complexity）は、それぞれ$N_G(b)b,N_D(b)b$で定義できます。これは、$1$回の反復で$b$個の確率的勾配を計算するからです。 \begin{align} N_G(b)b = \frac{A_G b^2}{(\epsilon_G^2 - C_G)b - B_G}, \text{ } N_D(b)b = \frac{A_D b^2}{(\epsilon_D^2 - C_D)b - B_D} \end{align} 先ほどと同様に微分して形状を調べてみましょう。 \begin{align} \frac{\mathrm{d}N_G(b)b}{\mathrm{d}b} = \frac{A_G b\{(\epsilon_G^2-C_G)b -2B_G\}}{\{(\epsilon_G^2-C_G)b-B_G\}^2} , \text{ } \frac{\mathrm{d}^2 N_G(b)b}{\mathrm{d}b^2} = \frac{2A_G B_G^2}{\{(\epsilon_G^2-C_G)b-B_G\}^3} \geq 0 \end{align} が成り立ちます。2階微分が常に$0$以上なので、$N_G(b)b$は$b$に対して凸関数であることが分かります。$N_D(b)b$についても同様です。 また、1階微分から、$N_G(b)b, N_D(b)b$を最小にする$b_G$と$b_D$は、次のように書けます。 \begin{align} b_G^\star := \frac{2B_G}{\epsilon_G^2-C_G}, \text{ } b_D^\star := \frac{2B_D}{\epsilon_D^2-C_D} \end{align} この$b_G^\star$と$b_D^\star$がクリティカルバッチサイズです。

ここまでで、GANの訓練にもクリティカルバッチサイズが存在することがわかりました。このクリティカルバッチサイズを、事前に知ることを目指します。
GANの目標は綺麗な生成画像を得ることです。学習の停止条件にも、生成画像の品質が利用されます。なので、ここでは生成器のクリティカルバッチサイズについて考えてみましょう。 生成器$G$のクリティカルバッチサイズは、次のように表すことができました。 \begin{align} b_G^\star := \frac{2B_G}{\epsilon_G^2 -C_G} \end{align} $B_G$や$C_G$の定義に立ち返って式変形をすると、$b_G^\star$の下界をoptimizerごとに次のように表すことができます。

• Adam

\begin{align} b_G^\star \geq \frac{\sigma_G^2}{\epsilon_G^3}\frac{\alpha^G}{(1-\beta_1^G)^3 \sqrt{\frac{\Theta}{1-\beta_2^G} \frac{1}{|S|^2}}} \end{align}

• AdaBelief

\begin{align} b_G^\star \geq \frac{\sigma_G^2}{\epsilon_G^3}\frac{\alpha^G}{(1-\beta_1^G)^3 \sqrt{\frac{4\Theta}{1-\beta_2^G} \frac{1}{|S|^2}}} \end{align}

• RMSProp

\begin{align} b_G^\star \geq \frac{\sigma_G^2}{\epsilon_G^3}\frac{\alpha^G}{\sqrt{\frac{\Theta}{|S|^2}}} \end{align} このとき、$\sigma_G^2 / \epsilon_G^3$のみ未知です。$\alpha^G, \beta_1^G, \beta_2^G$はoptimizerのハイパーパラメータなので、ユーザーが自由に設定できます。GANの訓練では、$\alpha^G=0.0001, \beta_1^G=0.5, \beta_2^G=0.999$などとするのが一般的です。 $\Theta$は生成器の次元でした。例えばDCGAN architectureならば、$\Theta=3,576,704$です。$|S|$はデータセットのデータの総数です。例えばLSUN Bedroomデータセットならば、$|S|=3,033,042$です。

それでは、未知数である$\sigma_G^2 / \epsilon_G^3$について考えてみましょう。これは決して事前には分からないので、クリティカルバッチサイズの測定値と、先ほどの下界の推定式を利用して逆算します。いくつかの準備が必要なので、順番に見ていきましょう。
　　1. $Nb-b$グラフ
DCGANを訓練して、LSUN Bedroomデータセットの実画像と瓜二つの生成画像を作ります。その学習のクリティカルバッチサイズの測定値を求めましょう。 クリティカルバッチサイズは、SFO計算量である$Nb$を最小にする$b$だったので、その測定値を知るためには、$Nb-b$グラフ（縦軸が$Nb$で、横軸が$b$のグラフ）を用意する必要があります。ここで、$N$は学習が収束するまでに必要なステップ数です。$b$はバッチサイズで、機械学習では慣例的に$2$の累乗を利用することが一般的です。理論的には$Nb$は$b$に対して凸関数であるはずなので、例えば次のようなグラフが得られるはずです。
　　2. FID
一般的に、GANの学習の停止条件にはFID（Fr\'echet Inception Distance）と呼ばれる指標が利用されます。FIDは2つの画像（※実際には2つの正規分布）の離れ具合を測る指標で、その値が低ければ低いほど2つの画像が似ていることを意味します。全く同じ画像同士のFIDは$0$となります。GANの学習が収束するとき、生成器は綺麗な生成画像を出力できるようになっているはずですから、その生成画像と実画像のFIDは十分に低いはずです。 ここまでのことを踏まえて、『学習が収束するまでに必要なステップ数$N$』を、『十分に低いFIDを達成するまでに必要なステップ数$N$』と読み替えます。 LSUN Bedroomデータセットを使って、DCGANを訓練する場合、FIDの極限は41.8程度です。学習の不安定性を考慮して、今回の実験では、FID=70を十分に低いFIDだとして、FID=70を停止条件にします。

　　3. 学習率
学習率は学習の結果に大きく影響するので、事前に適切な学習率を探索することが極めて重要です。GANの学習は非常に不安定です。DCGANは標準的なGANなので、特に学習率などのハイパーパラメータにとても敏感に影響を受けます。GANには生成器と識別器があるので、ある一定のステップ数の学習で、最も低いFIDを達成することができる学習率の組み合わせ$(\alpha^D, \alpha^G)$を探します。DCGANの場合は、次のような結果になります。

バッチサイズは$64$で固定し、Adamは$60000$steps、AdaBeliefは$120000$steps、RMSPropは$180000$stepsの結果です。青色が濃いほど低いFIDを達成できたことを意味します。 これによると、最も良い学習率の組み合わせは、optimizerごとに次のようになります。

なお、$\beta_1$と$\beta_2$は一般的な値を使用します。

それでは、いよいよバッチサイズを変えて、『FID=70を達成するまでに必要なステップ数』を計測していきましょう。 バッチサイズは大きければ大きいほどGPUをたくさん使用するので、計算機の都合上、今回の実験では、$b=[4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]$の範囲で実験しました。 $Nb-b$グラフを作ると、次のようになります。

予想通り、全てのoptimizerで凸関数になっています。このグラフの最小値をとるような$b$が、それぞれのクリティカルバッチサイズです。これで、Adamのクリティカルバッチサイズの測定値は$2^5=32$であることが分かりました。これを使って、未知数$\sigma_G^2 / \epsilon_G^3$を逆算しましょう。Adamの推定式をもう一度掲載します。 \begin{align} b_G^\star \geq \frac{\sigma_G^2}{\epsilon_G^3}\frac{\alpha^G}{(1-\beta_1^G)^3 \sqrt{\frac{\Theta}{1-\beta_2^G} \frac{1}{|S|^2}}} \end{align} これを未知数$\sigma_G^2 / \epsilon_G^3$について書き直して、既知のパラメータに使用した値を代入すると、 \begin{align} \frac{\sigma_G^2}{\epsilon_G^3} \leq 788.7 \end{align} とできます。これを使えば、AdaBeliefとRMSPropの推定値を計算できます。推定値は四角で、測定値は丸でマーキングしてあります。 さらに、生成器のモデルにDCGAN architectureを採用している場合は、この比$\sigma_G^2 / \epsilon_G^3$を適用できるので、別のGANでDCGAN architectureを採用している場合にもこの推定式は有効です。実際、WGAN-GPでCelebAデータセットを訓練する場合にも、推定値と測定値は近くなります。

DCGANがSection4.1で、WGAN-GPがSection4.2に当たります。RMSProp以外では完全に推定に成功していることが分かります。RMSPropで推定が上手くいかない原因は、RMSPropのクリティカルバッチサイズの推定式に$\beta_1$と$\beta_2$が含まれていないことであると考えられます。

Naoki Sato, Hideaki Iiduka: Existence and Estimation of Critical Batch Size for Training Generative Adversarial Networks with Two Time-Scale Update Rule, Proceedings of The 40th International Conference on Machine Learning, PMLR 202: ??–?? (2023)

Hiroki Naganuma, Hideaki Iiduka: Conjugate Gradient Method for Generative Adversarial Networks, Proceedings of The 26th International Conference on Artificial Intelligence and Statistics, PMLR 206: 4381–4408 (2023)