KL divergence loss goes to zero while training VAE I'm trying to train a variational autoencoder to perform unsupervised classification of astronomical images (they are of size 63x63 pixels). I'm using an encoder with 2 convolutional layers and a dense layer, and a similar structure for the decoder. I'm performing Xavier initialization of gradients. I'm using the Adam optimizer with a learning rate of 1e-4.
I observe that the KL divergence starts at very small values (roughly of the order of 1e-4) and suddenly vanishes after a few epochs while training, while my reconstruction loss reduces normally (I use MSE as the reconstruction loss). What could be the reason? Should I perform scaling of the losses separately?
This is my model.
initializer = glorot_normal()

def sampling(inputs):
  z_mean, z_log_var = inputs
  epsilon = k.random_normal(shape=(k.shape(z_mean)[0], 2), mean=0., stddev=0.1)
  return z_mean + k.exp(z_log_var) * epsilon

input_images = Input(shape = (63,63,1))
conv1 = Conv2D(16, (3,3), activation = 'relu')(input_images)
conv2 = Conv2D(8, (3,3), activation = 'relu')(conv1)
flattened = Flatten()(conv2)
x = Dense(4, activation = 'relu')(flattened)

z_mean = Dense(2, name = &quot;z_mean&quot;)(x)
z_log_var = Dense(2, name = &quot;z_log_var&quot;)(x)
z = Lambda(sampling, output_shape = (2,))([z_mean, z_log_var])

encoder = Model(input_images, [z_mean, z_log_var, z], name = &quot;encoder&quot;)

latent_inputs = Input(shape = (2,))
x = Dense(59*59*8, activation = 'relu')(latent_inputs)
x = Reshape((59,59,8))(x)
conv4 = Conv2DTranspose(8, (3,3), activation = 'relu')(x)
decoded = Conv2DTranspose(1, (3,3), activation = 'softmax')(conv4)

decoder = Model(latent_inputs, decoded, name = &quot;decoder&quot;)

z_mean, z_log_var, z = encoder(input_images)
vae_decoder_output = decoder(z)
vae = Model(input_images, vae_decoder_output, name = &quot;VAE&quot;)
vae.summary()

This is the loss function that I'm trying to implement.
recon = MSE(input_images, vae_decoder_output)
recon = k.mean(recon)
kl_loss = 1 + z_log_var - k.square(z_mean) - k.exp(z_log_var)
kl_loss = k.sum(kl_loss, axis = -1)
kl_loss *= -0.5
vae_loss = k.mean(kl_loss*10**3+recon)
vae.add_loss(vae_loss)
take a look at the paper &quot;Generating Sentences from a Continuous Space&quot; by Bowman. In Section 3.1 it is explained why LSTM_VAE tend to this behaviour:
&quot;This problematic tendency in learning is compounded by the lstm decoder’s sensitivity to subtle variation in the hidden states, such as that introduced by the posterior sampling process. This
causes the model to initially learn to ignore ~z and
go after low hanging fruit, explaining the data with
the more easily optimized decoder. Once this has
happened, the decoder ignores the encoder and little to no gradient signal passes between the two,
yielding an undesirable stable equilibrium with the
kl cost term at zero. We propose two techniques
to mitigate this issue.&quot;
....