Adding kld loss to my VAE-like model completely wrecks the performance

hubertrybka: <p>Hi,<br>
My model is an autoencoder-like neural net which takes a sparse bit vector of fixed len as an input and passes it through a VAE encoder. Reparametrization trick is used to sample the latent space and the latent vector is passed to GRU, which decodes it into a sequence of letters (tokens). The model is supposed to ‘describe’ the input bit vector with a sequence of tokens.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/3/b3e1e120e24d5c8edfd7c8162b4592a6cb308f62.png" data-download-href="https://discuss.pytorch.org/uploads/default/b3e1e120e24d5c8edfd7c8162b4592a6cb308f62" title="diagram(1)"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/3/b3e1e120e24d5c8edfd7c8162b4592a6cb308f62_2_690x471.png" alt="diagram(1)" data-base62-sha1="pFjpwedF7o5QAdXaD7zQmbq3MRA" width="690" height="471" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/3/b3e1e120e24d5c8edfd7c8162b4592a6cb308f62_2_690x471.png, https://discuss.pytorch.org/uploads/default/original/3X/b/3/b3e1e120e24d5c8edfd7c8162b4592a6cb308f62.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/b/3/b3e1e120e24d5c8edfd7c8162b4592a6cb308f62.png 2x" data-dominant-color="C0BFC1"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">diagram(1)</span><span class="informations">881×602 26.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
Now, the issue is as follows: when I trained the neural net, caring only about Cross-entropy loss (losss calculated on the output sequence) the model learns quickly and gets pretty good at the task (see graphs in my comment below). However, I need the encoder latent space to have normal distribution in order to sample from it and use the whole network as a generative model.</p>
<p><strong>If I include KLD loss in the backward pass, the optimizer seems to completely ingore the task of minimizing the cross-entropy and drives the KLD loss to almost zero. This happens even if KLD loss is multiplied by a very low weight prior to .backward()</strong></p>
<p>Thanks for all ideas and feedback.</p>
<pre><code class="lang-auto"># Training loop

    optimizer = torch.optim.Adam(model.parameters(), lr=learn_rate)
    for epoch in epochs_range:
        model.train()
        epoch_loss = 0
        kld_loss = 0
        for X, y in train_loader:
            X = X.to(device)
            y = y.to(device)
            optimizer.zero_grad()
            output, kld_loss = model(X, y, teacher_forcing=True, reinforcement=False)
            kld_loss = kld_loss * kld_weight
            loss = criterion(y, output)
            if kld_backward:
                (loss + kld_loss).backward()
            else:
                loss.backward()
            optimizer.step()

# KLD loss calculation inside model.encoder instance

    def kld_loss(mu, logvar):
        KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        return KLD

</code></pre>

hubertrybka: <p>I include the training metrics:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/1/a1d8fb3b1e1c0ffbc45afa35d84f9f364af3ef66.png" data-download-href="https://discuss.pytorch.org/uploads/default/a1d8fb3b1e1c0ffbc45afa35d84f9f364af3ef66" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/1/a1d8fb3b1e1c0ffbc45afa35d84f9f364af3ef66_2_690x294.png" alt="image" data-base62-sha1="n5LKHf8sActO2oMDXVdWmECdpPM" width="690" height="294" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/1/a1d8fb3b1e1c0ffbc45afa35d84f9f364af3ef66_2_690x294.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/1/a1d8fb3b1e1c0ffbc45afa35d84f9f364af3ef66_2_1035x441.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/1/a1d8fb3b1e1c0ffbc45afa35d84f9f364af3ef66_2_1380x588.png 2x" data-dominant-color="FDFDFD"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1394×595 35.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

hubertrybka: <p>FIXED: Here’s what I suspect happened: as the backward pass was called after each <strong>minibatch</strong>, and the logs are saved after each <strong>epoch</strong>, I did not realize what happened to the KLD and Cross-Entropy losses during the first epoch of training. What the issue turned out to be: before the first logs were collected, the model’s latent space was already highly organised with KLD loss approaching zero.<br>
The solution appears to be applying KLD annealing method during the first 10-20 epochs of training. I created a <a href="https://github.com/hubertrybka/vae-annealing" rel="noopener nofollow ugc">simple github repo</a> in order to make the technique easier to implement in different kinds of networks.<br>
Cheers.</p>