跳转至

Ed-VAE1

概要

代码本崩溃是训练使用离散表示空间的深度生成模型(如向量量化变分自编码器,VQ-VAE)时的常见问题. 作者们发现, 这种问题同样存在于另一种离散变分自编码器(dVAE)设计中——这类编码器通过直接学习代码本嵌入的分布来表示数据. 他们提出假设: 使用softmax函数获取概率分布会导致代码本崩溃, 因为该函数会为最佳匹配的代码本元素分配过度自信的概率. 本文创新性地引入证据深度学习(EDL), 通过分层贝叶斯建模替代softmax函数, 以解决dVAE的代码本崩溃问题. 与softmax方法不同, 我们通过证据理论框架来监控代码本嵌入概率分布获取过程的可信度. 在多个数据集上的实验表明, 所提出的EdVAE模型不仅有效缓解了代码本崩溃现象, 还提升了重建性能, 同时相较于dVAE和基于VQ-VAE的模型展现出更优的代码本利用率. 相关代码已开源在https://github.com/ituvisionlab/EdVAE.

dVAE

dVAE是和VQ-VAE在架构和目的上都引入了"离散"的概念, 从而在某些方面具有相似之处, 但是存在关键差异: dVAE的编码器直接输出离散的索引的分布(代码本嵌入的分布), VQ-VAE的编码器输出的是连续的向量, 然后通过量化操作后映射到离散码本中的向量.

由于输出one-hot向量是不可微的, 所以dVAE直接输出的是离散索引的分布, 或者说, 输出的是每个离散索引被选中的概率, 这个可以由Gumbel-Softmax实现.

动机

  • VQ-VAE面临"码本坍塌"问题: VQ-VAE通过确定性量化方式选择最近的潜在嵌入, 导致一部分码本嵌入未被充分使用, 限制了生成模型的表达能力
  • dVAE尽管引入随机性, 但仍存在"码本坍塌"问题: 作者发现, 即使dVAE的编码器学习码本嵌入上的分布(通过Categorical分布建模), 也仍然遭受码本坍塌的困扰.
  • Softmax函数可能是dVAE码本坍塌的诱因: 作者假设, dVAE中softmax函数获取码本嵌入概率质量的方式, 可能导致预测过度自信, 从而引发码本坍塌. 这指出了Softmax在某些情况下可能并不理想, 并暗示了探索替代方案的必要性.

贡献

  • 引入了dVAE的一个原创扩展, 这是一个分层的贝叶斯模型, 通过结合EDL(Evidential Deep Learning)使用了Dirichlet-Categorical分布.
  • 报告了由dVAE中使用的softmax概率导致的确认偏差问题的证据.
  • 观察到EDL的整合改善了dVAE的码本使用, 这通过困惑度(perplexity)来衡量.
  • 为dVAE在各种数据集上设定了基线结果, 并在重建性能, 码本使用和图像生成性能方面, 在大多数设置中超越了基线测量(他们的模型称为Evidential dVAE (EdVAE)).
  • 将Evidential dVAE与最先进的基于VQ-VAE的模型进行了比较, 证明了其模型在各种实验设置下性能接近或优于基于VQ-VAE的方法.

相关工作

dVAE

在离散变分自编码器(Discrete VAEs)中, 为了建模高维数据并学习其低维, 离散的潜在表示 \(z\), 我们通常将潜在变量 \(z\) 的先验分布 \(p(z)\) 和由编码器学习得到的近似后验分布 \(q(z|x)\) 都设定为类别分布(Categorical Distribution). 由于从离散的类别分布中进行采样是一个不可导的操作, 直接进行端到端训练会遇到梯度流中断的问题; 因此, 为了实现反向传播, 通常会采用Gumbel-Softmax技巧, 通过引入Gumbel噪声并使用其连续可导的"软化"版本来近似离散采样过程, 从而使得基于梯度的优化成为可能.

层次贝叶斯模型

层次贝叶斯模型被定义为遵循层次结构的一系列随机变量. 每个层次表示数据中不同程度的抽样或者变异性. 假设一个包含\(L\)个观测值的数据集\(X\)是由两组局部潜变量\(U\)\(Z\)生成的, 它们是独立的.


  1. Baykal, G., Kandemir, M., & Unal, G. (2024). EdVAE: Mitigating codebook collapse with evidential discrete variational autoencoders (No. arXiv:2310.05718). arXiv. https://doi.org/10.48550/arXiv.2310.05718