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), 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) vae.add_loss(vae_loss)