变分自动编码器中如何权衡KLD损耗与重构损耗

26

在我见过的VAE的几乎所有代码示例中，损失函数的定义如下（这是张量流代码，但是我看到theano，torch等类似。它也适用于卷积网络，但这也不太相关），仅影响轴的总和）：

# latent space loss. KL divergence between latent space distribution and unit gaussian, for each batch.
# first half of eq 10. in https://arxiv.org/abs/1312.6114
kl_loss = -0.5 * tf.reduce_sum(1 + log_sigma_sq - tf.square(mu) - tf.exp(log_sigma_sq), axis=1)

# reconstruction error, using pixel-wise L2 loss, for each batch
rec_loss = tf.reduce_sum(tf.squared_difference(y, x), axis=[1,2,3])

# or binary cross entropy (assuming 0...1 values)
y = tf.clip_by_value(y, 1e-8, 1-1e-8) # prevent nan on log(0)
rec_loss = -tf.reduce_sum(x * tf.log(y) + (1-x) * tf.log(1-y), axis=[1,2,3])

# sum the two and average over batches
loss = tf.reduce_mean(kl_loss + rec_loss)

但是，kl_loss和rec_loss的数值范围分别非常取决于潜在空间的暗淡和输入特征的大小（例如，像素分辨率）。用reduce_mean替换reduce_sum以获得每个z-dim KLD和每个像素（或特征）LSE或BCE是否明智？更重要的是，在对最终损失进行求和时，我们如何权衡潜在损失和重建损失？只是反复试验吗？还是有一些理论（或至少是经验法则）？我在任何地方都找不到任何相关信息（包括原始论文）。

我遇到的问题是，如果我的输入特征（x）尺寸和潜在空间（z）尺寸之间的平衡不是“最佳”，那么我的重构都非常好，但是学习到的潜在空间是非结构化的（如果x尺寸）很大，重构误差在KLD上占主导地位，反之亦然（重构不好，但如果KLD占优势，则学习的潜伏空间结构良好）。

我发现自己必须归一化重构损失（除以输入特征尺寸）和KLD（除以z尺寸），然后用任意权重因子手动加权KLD项（归一化是为了使我可以使用相同或相似的重量，与x或z的尺寸无关）。根据经验，我发现大约0.1可以在重构和结构化的潜在空间之间找到良好的平衡，这对我来说就像一个“最佳位置”。我正在寻找这方面的先前工作。

根据要求，使用上述数学符号（着眼于因重建错误而导致的L2损失）

{大号}_{升 一种 Ť Ë ñ Ť}^{（ 一世 ）} = - \frac{1个}{2} \sum_{Ĵ = 1个}^{Ĵ} （ 1个 + 日志 （ σ_{Ĵ}^{（ 一世 ）} ）^{2} - （ μ_{Ĵ}^{（ 一世 ）} ）^{2} - （ σ_{Ĵ}^{（ 一世 ）} ）^{2} ）

$\mathcal{L}_{latent}^{(i)} = -\frac{1}{2} \sum_{j=1}^{J}(1+\log (\sigma_j^{(i)})^2 - (\mu_j^{(i)})^2 - (\sigma_j^{(i)})^2)$

{大号}_{[R Ë C Ø ñ}^{（ 一世 ）} = - \sum_{ķ = 1个}^{ķ} （ ÿ_{ķ}^{（ 一世 ）} - X_{ķ}^{（ 一世 ）} ）^{2}

$\mathcal{L}_{recon}^{(i)} = -\sum_{k=1}^{K}(y_k^{(i)}-x_k^{(i)})^2$

{大号}^{（ 米 ）} = \frac{1个}{中号} \sum_{一世 = 1个}^{中号} （ {大号}_{升 一种 Ť Ë ñ Ť}^{（ 一世 ）} + {大号}_{[R Ë C Ø ñ}^{（ 一世 ）} ）

$\mathcal{L}^{(m)} = \frac{1}{M}\sum_{i=1}^{M}(\mathcal{L}_{latent}^{(i)} + \mathcal{L}_{recon}^{(i)})$

其中是潜矢量的维数（以及相应的均值和方差），是输入特征的维数，是最小批处理大小，上标表示第个数据点和是第个小批量的损失。 $J$ $z$ $\mu$ $\sigma^2$ $K$ $M$ $(i)$ $i$ $\mathcal{L}^{(m)}$ $m$

— 备忘录
source

17

对于在这篇文章上绊脚石并寻找答案的任何人，此twitter线程都增加了很多非常有用的见解。

即：

beta-VAE：使用受约束的可变框架学习基本的视觉概念

通过一些实验讨论我的确切问题。有趣的是，它们的（与我的归一化KLD权重相似）似乎也位于0.1的中心，较高的值提供了更多的结构化潜在空间，但重建效果较差，较低的值提供了更好的重建效果而结构化程度较低潜在空间（尽管它们的重点是专门学习纠缠的表示形式）。 $\beta_{norm}$

和相关阅读（讨论类似问题的地方）

— 备忘录
source

7

我想再增加一份与此问题相关的论文（由于目前信誉不佳，我无法发表评论）。

在本文的第3.1小节中，作者指定他们未能训练直接加权的可能性和KL差异对VAE的直接实施。在他们的情况下，尽管期望损失很小，但KL损失却不希望地减少到零。为了克服这个问题，他们建议使用“ KL成本退火”，这将KL散度项（蓝色曲线）的权重系数从0逐渐增加到1。

此变通办法也适用于Ladder VAE。

纸：

Bowman，SR，Vilnis，L.，Vinyals，O.，Dai，AM，Jozefowicz，R.和Bengio，S.，2015。从连续空间生成句子。arXiv预印本arXiv：1511.06349。

— ong
source