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

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)

take a look at the paper "Generating Sentences from a Continuous Space" by Bowman. In Section 3.1 it is explained why LSTM_VAE tend to this behaviour:

"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."



Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.