VAE重建损失函数

Question

我从TensorFlow文档给出的例子中学习，https://www.tensorflow.org/tutorials/generative/cvae#define_the_loss_function_and_the_optimizer:

VAEs通过最大化证据下界(ELBO)来训练边际对数似。

在实践中，优化单样本蒙特卡罗估计期望:logp(x|z) + logp(z) - logq(z|x).

损失函数实现为:

def log_normal_pdf(sample, mean, logvar, raxis=1):
log2pi = tf.math.log(2. * np.pi)
return tf.reduce_sum(
-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi),
axis=raxis)

def compute_loss(model, x):
mean, logvar = model.encode(x)
z = model.reparameterize(mean, logvar)
x_logit = model.decode(z)
cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=x_logit, labels=x)
logpx_z = -tf.reduce_sum(cross_ent, axis=[1, 2, 3])
logpz = log_normal_pdf(z, 0., 0.)
logqz_x = log_normal_pdf(z, mean, logvar)
return -tf.reduce_mean(logpx_z + logpz - logqz_x)

由于本例使用的是MINIST数据集，所以x可以归一化为[0,1]，这里使用的是sigmoid_cross_entropy_with_logits。

我的问题是:

1. 如果使用x > 1，可以使用什么样的损耗?
2. 我们是否可以使用其他损失函数作为VAE的重建损失，例如Huber损失(https://en.wikipedia.org/wiki/Huber_loss)?
3. 另一个使用MSE损失的例子(如下)，MSE损失是有效的ELBO损失来测量p(x|z)吗?

https://www.tensorflow.org/guide/keras/custom_layers_and_models putting_it_all_together_an_end-to-end_example

# Iterate over the batches of the dataset.
for step, x_batch_train in enumerate(train_dataset):
with tf.GradientTape() as tape:
reconstructed = vae(x_batch_train)
# Compute reconstruction loss
loss = mse_loss_fn(x_batch_train, reconstructed)
loss += sum(vae.losses)  # Add KLD regularization loss

Answer 1

在变分自编码器的损失函数中，联合优化两项:

1. 预测和标签之间的重建损失，就像在一个普通的自编码器
2. 参数化概率分布与假设真实概率分布之间的距离。在实践中，通常假设真实分布为高斯分布，距离以Kullback-Leibler散度
来测量。

对于重建损失部分，您可以选择适合您的数据的任何损失函数，包括MSE和Huber。通常来说，对输入特征进行规范化仍然是一个好主意。