[论文解读] Re-parameterizing VAEs for stability

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

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

一句话总结

本文提出了一种理论基础坚实的变分自编码器（VAEs）重参数化方案，通过解决正态分布中无界方差引起的数值不稳定性，实现了训练的稳定。通过用有界或线性-指数变体替代标准参数化方式——尤其针对编码和解码分布——该方法消除了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参数化也实现了训练稳定，三组不同随机种子下仅记录到一次高梯度事件，且无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次高梯度事件，表明不稳定性主要源于编码器方差，而非解码器方差。
• 所提出的参数化方式允许完全使用训练数据，消除了因跳过更新而丢弃宝贵梯度更新的启发式方法。