# 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 = "z_mean")(x)
z_log_var = Dense(2, name = "z_log_var")(x)
z = Lambda(sampling, output_shape = (2,))([z_mean, z_log_var])

encoder = Model(input_images, [z_mean, z_log_var, z], name = "encoder")

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 = "decoder")

z_mean, z_log_var, z = encoder(input_images)
vae_decoder_output = decoder(z)
vae = Model(input_images, vae_decoder_output, name = "VAE")
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)