[論文レビュー] Re-parameterizing VAEs for stability

David Dehaene, Rémy Brossard|arXiv (Cornell University)|Jun 25, 2021

Generative Adversarial Networks and Image Synthesis参考文献 12被引用数 5

ひとこと要約

この論文は、正規分布の無限大分散によって引き起こされる数値的不安定性に対処することで、訓練を安定化させる理論的裏付けのある再パrameterizationスキームを提案する。標準的なパrameterizationに代えて、特に符号化済みおよび復元済み分布のための有界または線形・指数関数的バリエーションを採用することで、NaN損失や高梯度を解消し、パフォーマンスに影響を与えることなく、深く複雑なVAEアーキテクチャの安定した訓練を可能にする。

ABSTRACT

We propose a theoretical approach towards the training numerical stability of Variational AutoEncoders (VAE). Our work is motivated by recent studies empowering VAEs to reach state of the art generative results on complex image datasets. These very deep VAE architectures, as well as VAEs using more complex output distributions, highlight a tendency to haphazardly produce high training gradients as well as NaN losses. The empirical fixes proposed to train them despite their limitations are neither fully theoretically grounded nor generally sufficient in practice. Building on this, we localize the source of the problem at the interface between the model's neural networks and their output probabilistic distributions. We explain a common source of instability stemming from an incautious formulation of the encoded Normal distribution's variance, and apply the same approach on other, less obvious sources. We show that by implementing small changes to the way we parameterize the Normal distributions on which they rely, VAEs can securely be trained.

研究の動機と目的

• 深層VAEにおける継続的な数値的不安定性、特に訓練中の高梯度およびNaN損失を解消すること。
• 標準的なVAEパラメータライゼーションにおける不安定性の根本的要因を特定すること、特に符号化された正規分布の分散における要因。
• 勾配クリッピングや更新スキップといった経験的修正に代わる理論的裏付けのある代替手法を提供すること。
• ヒューリスティックなデータフィルタリングや更新スキップに依存せずに、最新のVAEアーキテクチャ（例：非常に深層VAE（VDVAE））の安定した訓練を保証すること。
• 最小限のアーキテクチャ的変更で、パフォーマンスの低下を伴わずに安定した訓練が可能であることを示すこと。

提案手法

• 符号化された正規分布のためのUpBounded1パラメータライゼーションを導入し、分散が有界に保たれ、発散を回避することを保証する。
• 復元された分布のためのBoundedパラメータライゼーションを提案し、分散が入力データの範囲（例：画像生成における画素値）内に保たれることを保証する。
• 符号化された分散のためのExpLinパラメータライゼーションを適用し、ニューラルネットワークの出力を線形変換することで正の値と安定性を確保する。
• エンコーダーおよびデコーダーにおける正規分布の標準的再パラメータライゼーションを、これらの新しいパラメータライゼーションに置き換えることで、勾配フローの安定化を図る。
• ベースラインとしてデコーダー経路で一定分散（例：1.0）を用い、表現力と安定性をテストする。これによりモデルのパフォーマンスが制限されないことが示された。
• CIFAR10で訓練された非常に深層VAE（VDVAE）に対して、勾配が閾値を超えた際に更新をスキップするオリジナル実装と比較して、手法を検証する。

実験結果

リサーチクエスチョン

• RQ1 現代の深層VAEにおいて、正規分布の分散パラメータライゼーションに起因する数値的不安定性（例：NaN損失や極端な勾配）の原因は何か？
• RQ2 経験的ヒューリスティクスに代わる理論的裏付けのある再パラメータライゼーションスキームによって、VAE訓練における不安定性を解消できるか？
• RQ3 標準的な分散パラメータライゼーションを有界または線形・指数関数的代替に置き換えることで、訓練の安定性が向上し、モデルパフォーマンスが低下しないか？
• RQ4 高勾配による更新スキップなしに、VAEを信頼性高く訓練できるか？これにより訓練データの完全な利用が可能になるか？
• RQ5 特に改善されたエンコーダーのパラメータライゼーションと組み合わせた場合、デコーダーの分散を一定値に固定しても、VAEの表現力が保持されるか？

主な発見

• CIFAR10でVDVAEを訓練した際、符号化済み分散のためのUpBounded1パラメータライゼーションにより、NaN損失および高勾配イベントが完全に解消され、更新スキップがゼロに減少した。
• ExpLinパラメータライゼーションも訓練を安定化させ、3つのランダムシードで合計1件の高勾配イベントしか記録されず、NaNも発生しなかった。
• デコーダー分散を一定値（1.0）に固定してもパフォーマンスが低下せず、ExpLin/Cst1のテストNELBO値は2.006 ± 0.008と、オリジナルモデルと同等であった。
• UpBounded1を適用した変更版VDVAEは、テストNELBOが2.020 ± 0.006を達成し、オリジナルモデル（2.014 ± 0.009）と同等のパフォーマンスを示したが、完全に安定していた。
• UpBounded1パラメータライゼーションでは、一定分散ベースライン（UpBounded1/Cst1）で2,645件の高勾配イベントが記録されたことから、不安定性の主な原因はデコーダーの分散ではなく、エンコーダーの分散に起因することが示された。
• 提案されたパラメータライゼーションにより、更新スキップのヒューリスティクスに依存せず、訓練データの完全な利用が可能になった。