KLD loss goes NaN during VAE training

nk00: <p>Hello!</p>
<p>I’ve trained a stand-alone VAE based on the <a href="https://github.com/pytorch/examples/blob/master/vae/main.py" rel="nofollow noopener">PyTorch example</a> and a few other bits of code found on github - it works well and my output images look quite good. All of the examples dealt with MNIST but my model uses ImageNet images so it’s a big bigger than the examples. Apart from that, it doesn’t differ too much.</p>
<p>When I then want to use the VAE model in a bigger network where the output of the VAE is fed to a CNN and their losses are summed (there is a reason for that, I promise) the KLD loss of the VAE part is NaN. The VAE and CNN form one model but I return the output of the VAE (assigned to a separate variable before it even enters the CNN half) and calculate my VAE loss using the same function I used in the stand-alone VAE example.</p>
<p>From the very 1st epoch the KLD loss is NaN (BCE loss is fine). Is there a way to stabilise this behaviour? What could make the model work as a stand-alone system but fail when linked with anything else?</p>
<p>The KLD loss tends to be much much bigger in the 1st epoch so I thought it might be a data casting problem and added a scaling factor of 1e-10 but even that didn’t help. I considered setting KLD to some small value in the 1st epoch to prevent the overshoot but then I just get a NaN in the 2nd epoch…</p>
<p>This is my loss function:</p>
<pre><code class="lang-auto">def vae_loss(reconstructed_x, x, mean, logvar, batch_size):
    BCE = F.binary_cross_entropy(reconstructed_x, x, reduction='sum') 
    # see Appendix B from VAE paper:
    # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
    # https://arxiv.org/abs/1312.6114
    # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    KLD = -0.5 * torch.sum(1 + logvar - mean.pow(2) - logvar.exp())
    
    #normalise the outputs so divide by spatial dimensions + batch size
    BCE /= (3*model_size*model_size*batch_size)
    KLD /= (3*model_size*model_size*batch_size)
    total = BCE+KLD
    return total, BCE, KLD
</code></pre>
<p>Then, during training, I say</p>
<pre><code class="lang-auto">  output, ae_output, mean, logvar = model(input)
  cnn_loss = cnn_criterion(output, labels) # cross entropy loss
  ae_loss, bce, kld = vae_loss(ae_output, input, mean, logvar, batch_size)
  total_loss = ae_loss + cnn_loss
</code></pre>
<p>Has anyone dealt with something similar? Any ideas how to fix this?</p>

YingLaiLin: <p>I have met the same problem when using VAE for representation learning.<br>
To solve this problem, you must be know what lead to nan during the training process. I think the  logvar.exp() in the following fomula lead to overflow in the running process</p>
<pre><code class="lang-auto">KLD = -0.5 * torch.sum(1 + logvar - mean.pow(2) - logvar.exp())
</code></pre>
<p>so, we need to limit logvar in a specific range by some means.<br>
So, you can initialize weights of VAE in a small range, which is recommended in the range [-0.08, 0.08] to make sure the logvar is small, thus exp can not lead to overflow numerically.<br>
Plus, if you are using CNN for image classification, you need to make image zero-means and unit variance.</p>
<p>For details, you can visit this page. Best wishes.<br>
<a href="https://github.com/y0ast/VAE-Torch/issues/3" rel="nofollow noopener">Use of exp in the Reparametrization</a></p>

cherepanovic: <blockquote>
<p>To solve this problem, you must be know what lead to nan during the training process. I think the logvar.exp() in the following fomula lead to overflow in the running process</p>
</blockquote>
<p>could also be this part</p>
<pre><code class="lang-auto">
BCE = F.binary_cross_entropy(reconstructed_x, x, reduction='sum'
</code></pre>
<p>could you provide your exception?</p>

bsridatta: <aside class="quote no-group" data-username="YingLaiLin" data-post="2" data-topic="42305">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/yinglailin/48/12548_2.png" class="avatar"> YingLaiLin:</div>
<blockquote>
<p>o solve this problem, you must be know what lead to nan during the training process. I think the logvar.exp() in the following fomula lead to overflow in the running process</p>
</blockquote>
</aside>
<p>I can second this, I clip log var value and also just add an assertion to detect NaN values</p>