Can I use the mse loss function along with a sigmoid activation in my VAE?

Question

I have implemented a VAE with a mse loss for the reconstruction loss and a sigmoid activation in my last layer of the decoder.

For my use-case the reconstructed images seemed fair. Most examples on MNIST use the BCE as a loss, but since a paper on my use-case stated that they used the mse loss I also implemented it.

After some more research on the loss function, I got a little confused. I did not quite understand if I can use the mse loss along with a sigmoid activation in my last layer.

Can someone tell me if I can use the mse loss with the sigmoid activation?

My VAE has the following layer architecture:

latent_dim = 100 
encoder_inputs = keras.Input(shape=(32, 32, 3)) 
x = layers.Conv2D(32, 4, strides=2, padding="same")(encoder_inputs)   
x = layers.LeakyReLU(alpha=0.2)(x)
x = layers.Conv2D(32, 3, strides=1, padding="same")(x)   
x = layers.LeakyReLU(alpha=0.2)(x)
x = layers.Conv2D(64, 3,strides=1, padding="same")(x)
x = layers.LeakyReLU(alpha=0.2)(x)
x = layers.Conv2D(128, 4,strides=2, padding="same")(x)  
x = layers.LeakyReLU(alpha=0.2)(x)
x = layers.Conv2D(64, 3,strides=1, padding="same")(x)
x = layers.LeakyReLU(alpha=0.2)(x)
x = layers.Conv2D(32, 3,strides=1, padding="same")(x)
x = layers.LeakyReLU(alpha=0.2)(x)
x = layers.Conv2D(100, 8,strides=1, padding="valid")(x)   
x = layers.LeakyReLU(alpha=0.2)(x)
encoder_test = x
x = layers.Flatten()(x)
z_mean = layers.Dense(latent_dim, name="z_mean")(x)
z_log_var = layers.Dense(latent_dim, name="z_log_var")(x)
z = Sampling()([z_mean, z_log_var])
encoder = keras.Model(encoder_inputs, [z_mean, z_log_var, z], name="encoder")
encoder.summary()

latent_inputs = keras.Input(shape=(latent_dim,))
x = layers.Reshape((1, 1, latent_dim))(latent_inputs)
x = layers.Conv2DTranspose(latent_dim, 8, strides=1, padding="valid")(x)
x = layers.LeakyReLU(alpha=0.2)(x)
x = layers.Conv2DTranspose(32, 3, strides=1, padding="same")(x)
x = layers.LeakyReLU(alpha=0.2)(x)
x = layers.Conv2DTranspose(64, 3, strides=1, padding="same")(x)
x = layers.LeakyReLU(alpha=0.2)(x)
x = layers.Conv2DTranspose(128, 4, strides=2, padding="same")(x)
x = layers.LeakyReLU(alpha=0.2)(x)
x = layers.Conv2DTranspose(64, 3, strides=1, padding="same")(x)
x = layers.LeakyReLU(alpha=0.2)(x)
x = layers.Conv2DTranspose(32, 3, strides=1, padding="same")(x)
x = layers.LeakyReLU(alpha=0.2)(x)
x = layers.Conv2DTranspose(32, 4, strides=2, padding="same")(x)
x = layers.LeakyReLU(alpha=0.2)(x)

decoder_outputs = layers.Conv2DTranspose(3, 3, activation="sigmoid", padding="same")(x)
decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")
````

Answer 1

With the MSE loss, you are welcome to use any activation function in your last layer, including the sigmoid function. However, since you are using the sigmoid activation function anyway, then it doesn't hurt to try both the MSE loss and the BCE loss. This is because the BCE loss enables faster convergence.

Tips on training the VAE

• Since you are implementing a VAE, I would suggest you turn off any biases in your linear and convolutional layers, as this will avoid collapsing your reconstructions to the mean value of the input images. This is discussed in footnote 1 in section 3.3 (page 5) in the paper titled Deep One-Class Classification by Lukas Ruff et al.
• Also, consider keeping your batch size small at the beginning of training, to around 8 images per batch. I've observed that a smaller batch size greatly helps training.