「ハイパーパラメータ選択にロバストなVAE」に関する研究がICLR2023 notable-top-5%で採択されました

こんにちは、人工知能研究所長谷川です。このたび、富士通研究所とトロント大学との共同研究で得られた「ハイパーパラメータ選択にロバストなVAE」に関する研究論文が、機械学習の主要な国際会議であるICLR2023にnotable-top-5% (従来のOral)として採択されましたので、その内容を紹介いたします。

対象論文

タイトル: Multi-Rate VAE: Train Once, Get the Full Rate-Distortion Curve (Juhan Bae, Michael R. Zhang, Michael Ruan, Eric Wang, So Hasegawa, Jimmy Ba, Roger Baker Grosse)
発表会議: The Eleventh International Conference on Learning Representations (ICLR 2023) openreview.net

採択された論文の内容

Introduction

昨今何かと話題に上るGenerative AIですが、それらGenerative AIは生成モデルという技術を核としています。Variational AutoEncoder (VAE) [Kingma, et al] は数ある生成モデルの中でも主要モデルの一つです。VAEはAutoEncoderのようにエンコーダとデコーダによって構成されます。潜在変数が何らかの事前確率分布に沿っていると仮定し、エンコーダの事後分布が事前分布に沿うように学習を進めます。富士通研究所ではこれまでにVAEの定量的解析を行いICML2020 [Kato, et al] とICML2021 [Nakagawa, et al]で研究成果を発表しました。

VAEの一般化である $\beta$ -VAE [Higgins, et al]の目的関数は再構成誤差(distortion)と事後分布と事前分布のKullback-Leibler divergence(rate)という二つの項によって構成されており、定式化すると以下のようになります。

式の中で $\beta$ はKL divergence項に係る重みであり、distortionとrateのバランスを調整する役割を果たします。 $\beta$ の値を変更しながら学習を行った結果、図1のようなレート歪曲線が生まれます。

図1から分かるようにrateとdistortionはトレードオフの関係にあり、主に2つの領域が存在します。

Low distortion, high rate: 学習されたモデルは高い品質でデータを再構成出来ますが、不自然なデータをサンプルします。
Low rate, high distortion: 学習されたモデルは事前分布に沿って多様性を持ったデータをサンプル出来ますが、データの再構成は不十分です。

最適なrateとdistortionを保持するモデルを得るためには、 $\beta$ の値を変えながら複数回 $\beta$ -VAEを学習し、タスクに応じてそれらのモデルの推論結果を比較し選択する必要があります。簡素なニューラルネットワークなら複数回の学習は大したオーバーヘッドにはなりませんが、昨今の潮流である大規模なニューラルネットワークなら計算コストは莫大になります。そこで、我々は以下のような疑問を抱きました。たった一回の学習でレート歪曲線を得られるか？今回ICLR2023に採択された論文ではこの疑問に答える手法を提案しました。本手法によって $\beta$ をチューニングするコストを抑えられます。また、本手法が様々な評価指標において $\beta$ の選択に関わらずロバストであることも判明しました。

Method

始めに手法の要である応答関数を紹介します。応答関数 (Response function) はハイパーパラメータの値を、そのパラメータで学習した際に得られる最適なモデルパラメータに変換する関数です。VAEの場合応答関数は $\beta$ を、 $\beta$ -VAEの目的関数を最小化するエンコーダ・デコーダの最適パラメータへと変換します。我々の提案するMulti-rate VAE (MR-VAE)は図2のように、VAEの各層の重みとバイアスを単にスケールするだけで応答関数を構成します。

より具体的には、 VAEの $i$ 番目の層において重みとバイアスを $\rm{W}^{(i)} \in \mathbb{R}^{m_{i+1} \times m_i}$ と $\rm{b}^{(i)} \in \mathbb{R}^{ m_{i+1} }$ として表し、応答関数を構成するハイパーネットワーク $\psi$ を以下のように設定します。

ここで $\odot$ と $\odot_{\rm{row}}$ は要素毎の積と列毎のスケーリングをそれぞれ表しています。また、活性化関数 $\sigma^{(i)}$ を以下のように定義します。

MR-VAEは各層に $2m_{i}+1$ 分のパラメータと2回の要素和を順伝搬において追加で必要とするだけなので、計算コストも非常に軽いです。また、我々はこの応答関数が線形VAEの応答関数と一致していることも証明しています。詳しくは論文の付録をご覧ください。

全ての $\beta$ に対して最適なモデルパラメータを学習するため、我々は目的関数を以下のように提案しました。

$\mathcal{U} [\rm{log}(a), \rm{log}(b)$ ] は $\rm{log}(a)$ と $\rm{log}(b)$ を範囲とした一様分布であり、Self-Tuning Networkで用いられる目的関数と類似しています [Lorraine, et al][MacKay, et al][Bae, et al]。この目的関数によって、応答関数は $a$ から $b$ の全ての範囲の $\beta$ において最適なモデルパラメータを得るように学習されます。

Experiments

MR-VAEに関して我々は以下4つの疑問に答える実験を行いました。

MR-VAEは一度の学習でレート歪曲線を描けるか？
モダンサイズのネットワークに適用可能か？
MR-VAEは $\beta$ のサンプリング範囲に過敏か？
MR-VAEは他のVAEへと適用可能か？

1. MR-VAEは一度の学習でrate-distortion curveを描けるか？

まず、線形VAEにおいて我々の理論の妥当性を検証します。 $\beta$ を0.01から10.0の間でサンプリングし、MNISTにおいてMR-VAEを一度だけ学習しました。比較のため、線形VAEを $\beta$ の値を変えながら10回学習しました。得られたレート歪曲線を図3に示します。図3から明らかなように、MR-VAEによって得られたカーブが10回の線形VAEの学習によって得られたカーブに沿っています。

2. モダンサイズのネットワークに適用可能か？

次に、我々の手法がResNet等のよく用いられるネットワークでも通用するのか検証します。畳み込みニューラルネットワーク(Convolutional Neural Network, CNN)とResNetベースのネットワークをMNIST, Omniglot, CIFAR-10, SVHN, そしてCelebAの5つのデータセットを用いてそれぞれ学習しました。また、階層型VAEの代表的手法であるNVAE [Vahdat, et al] をMNIST, Omniglot, CelebAの3つのデータセットに対して学習しました。最後に、自己回帰LSTM VAEをYahooデータセットを用いて [He, et al] に沿って学習しました。これらの学習によって得られた各レート歪曲線を図4に示します。様々なアーキテクチャ、データセットに対してレート歪曲線が得られています。

また、MR-VAEと通常のVAEとの比較のため、rateやdistortionだけでなく他の評価指標と $\beta$ の関係にも着目しました。Fréchet Inception Distance [Heusel, et al] とActive Units [Burda, et al] を幾つかのデータセットにおいて計算し、結果を図5に示します。興味深いことに、MR-VAEは通常のVAEに比べて高い $\beta$ において意味のある潜在変数が得られ (高いActive Units)、低い $\beta$ において自然な画像を生成できています (低いFID)。これは、各 $\beta$ においてエンコーダ・デコーダのパラメータを共有しているため、様々な $\beta$ で学習された知識が転用されているためだと考えられます。このようにMR-VAEはActive UnitsとFIDに対して従来のVAEと比べて頑健であるという非常に面白い結果が得られました。

3. MR-VAEは $\beta$ のサンプリング範囲に過敏か？

MR-VAEは二つのハイパーパラメータ $a$ と $b$ を用いています。これらは $\beta$ のサンプリング範囲を決めるパラメータです。これらのハイパーパラメータの影響を検証します。ResNetベースのネットワークをOmniglotデータセットを用いて学習した結果を図6に示します。左図は、 $b$ = 10に固定して $a$ の値を0.001, 0.01, 0.1, 1.0と変化させた場合です。右図は、 $a$ = 0.01に固定して $b$ の値を10, 1.0, 0.1へと変化させた場合です。サンプリング範囲が異なるため、レート歪曲線の長さに変化はありますが、それぞれのカーブは一つの大きなカーブ上に乗っていることが分かります。

4. MR-VAEは他のVAEへと適用可能か？

既にNVAEと自己回帰LSTM VAEについては検証していますが、再構成誤差とKL divergence以外の目的関数を持つVAEも検証します。本実験では $\beta$ -TCVAEを用います。 $\beta$ -TCVAEはKL divergenceではなくtotal correlationを目的関数として用いています。この場合、 $\beta$ は再構成誤差とtotal correlationをバランスする役割を担います。我々はdSpritesデータセットに対してMulti-layer Perceptron (MLP)で構成されたMR-VAEを学習し、Mutual Information Gap (MIG)上の性能で比較しました。結果を図7に示します。この結果も非常に興味深いことに、MR-VAEは $\beta$ に対して頑健である上に $\beta$ -TCVAEよりも高い性能を持つことが分かります。

Conclusion

本記事においてMulti-rate VAE (MR-VAE) を紹介しました。この手法によって一度の学習でレート歪曲線が得られます。シンプルで且つ計算コストも軽い応答関数を導入した点が本手法の要になります。実験によって、様々なモダンサイズのネットワークやVAEにおいても、MR-VAEが一度の学習でレート歪曲線を獲得出来ることを示しました。また、Active UnitsやFIDといった別の評価指標においてハイパーパラメータ $\beta$ の値に関わらず頑健であるという興味深い結果も得られました。

今後の展開としては、提案手法の他生成モデルへの応用が期待されます。生成品質を定量化することは困難であり、生成データはトレードオフの品質を保持します。例えば、超解像タスクにおいてはdistortionとperceptionはトレードオフです [Blau et al]。このような特性によって、ハイパーパラメータを変更して何度も生成モデルを学習する必要があります。昨今の生成モデルのようにモデルが巨大である場合には、月や年単位でチューニングが必要になるでしょう。我々は提案手法がそのような問題に対しての解決策の一つであると考えています。

Reference

Diederik P. Kingma, et al., “Auto-Encoding Variational Bayes”, International Conference on Learning Representations (ICLR), 2014
Keizo Kato, et al., “Rate-Distortion Optimization Guided Autoencoder for Isometric Embedding in Euclidean Latent Space”, International Conference on Machine Learning (ICML) 2020
Akira Nakagawa, et al., “Quantitative Understanding of VAE as a Non-linearly Scaled Isometric Embedding”, International Conference on Machine Learning (ICML) 2021
Irina Higgins, et al., “beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework”, International Conference on Learning Representations (ICLR), 2017
Arash Vahdat, et al., “NVAE: A Deep Hierarchical Variational Autoencoder”, Neural Information Processing Systems (NeurIPS), 2020
Junxian He, et al., “Lagging Inference Networks and Posterior Collapse in Variational Autoencoders”, International Conference on Learning Representations (ICLR), 2019
Martin Heusel, et al., “Gans trained by a two time-scale update rule converge to a local nash equilibrium”, Neural Information Processing Systems (NeurIPS), 2017
Yuri Burda, et al., “Importance Weighted Autoencoders”, arXiv preprint arXiv:1509.00519, 2015
Ricky T. Q. Chen, et al., “Isolating Sources of Disentanglement in VAEs”, Neural Information Processing Systems (NeurIPS), 2018
Yochai Blau, et al., “The Perception-Distortion Tradeoff”, Computer Vision and Pattern Recognition Conference (CVPR), 2018
Jonathan Lorraine, et al. “Stochastic hyperparameter optimization through hypernetworks”, arXiv preprint arXiv:1802.09419, 2018
Matthew MacKay, et al., “Self-tuning networks: Bilevel optimization of hyperparameters using structured best-response functions”, arXiv preprint arXiv:1903.03088, 2019
Juhan Bae, et al., “Delta-stn: Efficient bilevel optimization for neural networks using structured response jacobians”, Neural Information Processing Systems (NeurIPS), 2020