<h3>The example</h3>
<p>I'm familiar with that example, and I think the <strong>28x28</strong> multiplier is justified because of the operation <code>tf.reduce_mean(kl_loss)</code> which takes the average loss of all the pixels in the image which would result in a number between 0 and 1 and then multiplies it by the number of pixels.  Here's <a href="https://www.tensorflow.org/guide/keras/custom_layers_and_models" rel="nofollow noreferrer">another take</a> with an external training loop for creating a VAE.</p>
<h3>The problem is posterior collapse</h3>
<p>The above would not be an issue since it's just multiplication by a constant if not for as you point out the <em>KL divergence</em> term.  The KL loss acts as a regularizer that penalizes latent variable probability distributions that when sampled using a combination of Gaussians are different than the samples created by the encoder. Naturally, the question arises, how much should be reconstruction loss and how much should be the penalty.  This is an area of research.  Consider <em>β-VAE</em> which purportedly serves to disentangle representations by increasing the importance of KL-loss, on the other hand, increase <strong>β</strong> too much and you get a phenomenon known as posterior collapse <a href="https://arxiv.org/abs/1910.00698" rel="nofollow noreferrer">Re-balancing Variational Autoencoder Loss for Molecule Sequence Generation</a> limits <strong>β</strong> to 0.1 to avoid the problem.  But it may not even be that simple as explained in <a href="https://arxiv.org/abs/1912.10702" rel="nofollow noreferrer">The Usual Suspects? Reassessing Blame for VAE Posterior Collapse</a>. A thorough solution is proposed in <a href="https://arxiv.org/abs/1903.05789" rel="nofollow noreferrer">Diagnosing and Enhancing VAE Models</a>. While <a href="https://arxiv.org/abs/2002.07514" rel="nofollow noreferrer">Balancing reconstruction error and Kullback-Leibler divergence in Variational Autoencoders
</a> suggest that there is a more simple deterministic (and better) way.</p>
<h3>Experimentation and Extension</h3>
<p>For something simple like Minst, and that example, in particular, try experimenting.  Keep the 28x28 term, and arbitrarily multiply kl_loss by a constant <code>B</code> where 0 &lt;= B &lt; 28*28.  Follow the kl loss term and the reconstruction loss term during training and compare it to the first reference graphs.</p>
Should reconstruction loss be computed as sum or average over image for variational autoencoders?
<p>I am following this variational autoencoder tutorial: <a href="https://keras.io/examples/generative/vae/" rel="nofollow noreferrer">https://keras.io/examples/generative/vae/</a>.</p>
<p>I know VAE's loss function consists of the reconstruction loss that compares the original image and reconstruction, as well as the KL loss. However, I'm a bit confused about the reconstruction loss and whether it is over the entire image (sum of squared differences) or per pixel (average sum of squared differences). My understanding is that the reconstruction loss should be per pixel (MSE), but the example code I am following multiplies MSE by <code>28 x 28</code>, the MNIST image dimensions. Is that correct? Furthermore, my assumption is this would make the reconstruction loss term significantly larger than the KL loss and I'm not sure we want that.</p>
<p>I tried removing the multiplication by (28x28), but this resulted in extremely poor reconstructions. Essentially all the reconstructions looked the same regardless of the input. Can I use a lambda parameter to capture the tradeoff between kl divergence and reconstruction, or it that incorrect because the loss has a precise derivation (as opposed to just adding a regularization penalty).</p>
<pre><code>reconstruction_loss = tf.reduce_mean(
    keras.losses.binary_crossentropy(data, reconstruction)
)
reconstruction_loss *= 28 * 28
kl_loss = 1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)
kl_loss = tf.reduce_mean(kl_loss)
kl_loss *= -0.5
total_loss = reconstruction_loss + kl_loss
</code></pre>