[论文解读] Efficient-VDVAE: Less is more

Louay Hazami, Rayhane Mama+1 more|arXiv (Cornell University)|Mar 25, 2022

Seismic Imaging and Inversion Techniques被引用 10

一句话总结

本文提出 Efficient-VDVAE，即 Very Deep VAE 的改进版本，通过简单的架构调整，实现最多 2.6 倍的收敛速度提升、20 倍的显存减少，以及训练稳定性的改善。尽管经过这些优化，其在 7 个图像数据集上的负对数似然（NLL）性能仍达到或超过当前最先进水平，同时表明仅有 3% 的潜在维度编码了图像的大部分信息。

ABSTRACT

Hierarchical VAEs have emerged in recent years as a reliable option for maximum likelihood estimation. However, instability issues and demanding computational requirements have hindered research progress in the area. We present simple modifications to the Very Deep VAE to make it converge up to $2.6 imes$ faster, save up to $20 imes$ in memory load and improve stability during training. Despite these changes, our models achieve comparable or better negative log-likelihood performance than current state-of-the-art models on all $7$ commonly used image datasets we evaluated on. We also make an argument against using 5-bit benchmarks as a way to measure hierarchical VAE's performance due to undesirable biases caused by the 5-bit quantization. Additionally, we empirically demonstrate that roughly $3\%$ of the hierarchical VAE's latent space dimensions is sufficient to encode most of the image information, without loss of performance, opening up the doors to efficiently leverage the hierarchical VAEs' latent space in downstream tasks. We release our source code and models at https://github.com/Rayhane-mamah/Efficient-VDVAE .

研究动机与目标

• 解决 Very Deep VAE（VDVAE）在训练过程中存在的不稳定性和高计算成本问题。
• 在不牺牲模型性能的前提下，提升收敛速度和内存效率。
• 挑战将 5 位量化基准作为分层 VAE 评估标准的公平性。
• 探究分层 VAE 中潜在空间的有效维度及其在下游任务中的实用性。
• 实现分层 VAE 在表征学习和生成建模中更高效的部署。

提出的方法

• 引入一种改进的 VDVAE 架构，采用有界输出层以稳定训练并降低梯度方差。
• 移除解码器中的无界输出层，替换为有界激活函数，防止梯度发散。
• 采用基于先验的剪枝策略，使不活跃的潜在维度坍缩，从而减少显存占用。
• 使用基于阈值的方法识别并剪除平均 KL 散度较低的潜在维度，表明其信息含量较低。
• 采用分层 ELBO 目标函数并结合自顶向下的推理机制，在保持 VDVAE 核心结构的同时提升数值稳定性。
• 在 8 位全彩图像数据集上评估性能，以避免 5 位量化带来的偏差。

实验结果

研究问题

• RQ1 能否在不损害似然性能的前提下，显著提升 VDVAE 的训练稳定性和效率？
• RQ2 5 位量化对分层 VAE 的基准测试有何影响？它是否是一个可靠的评估标准？
• RQ3 在分层 VAE 中，潜在空间的多少比例实际上对准确的图像重建是必要的？
• RQ4 能否在不影响重建质量或似然分数的前提下，对潜在空间进行大幅剪枝？
• RQ5 在保持最大似然估计的最先进性能的同时，内存和计算成本最多可降低多少？

主要发现

• Efficient-VDVAE 相较于标准 VDVAE，实现最多 2.6 倍的收敛速度提升，显存负载减少最多 20 倍。
• 该模型在所有 7 个评估图像数据集上的负对数似然（NLL）性能与当前最先进模型相当或更优。
• 移除无界输出层可稳定训练过程，且在 8 位基准上无性能下降，同时保持 NLL 性能不变。
• 仅 3% 的潜在维度负责编码图像的大部分信息，经验证：剪枝 97% 的维度后，重建损失仍保持等效。
• 将潜在空间剪枝至 3% 时，编码数据集大小最多可减少 33 倍，同时在 FFHQ 256×256 数据集上保持近乎相同的 NLL 分数。
• 本文认为，5 位基准引入了量化伪影和偏差，因此不适合用于评估分层 VAE。