# Activation function in the mean and variance layer in VAE?

I have come across several different codes of vae just confused about should I apply activation function or even batch normalization to mean and variance layer?

For example, the most common is the following one without specifying the activation function in the mean layer.

However, I also looked at this following code, it applied batch norm and relu to it

Which is correct? Many thanks!

• Is there any reason to believe that the two programs have the same purpose?
– Sycorax
Commented Jan 12, 2020 at 0:44
• @SycoraxsaysReinstateMonica, actually I am not clear about that, I would assume the purpose is training a valid vae. I am just confused about the second code where it applies the "Relu" to the mean layer which I have not seen elsewhere. Is that correct? Commented Jan 12, 2020 at 0:49
• Best to ask the author of the code what problem they were trying to solve. All we can do is guess. Note that when the VAE is parameterized as a normal distribution, the mean can be any real number and the log variance can be any real number. ReLU activations are non-negative, so this implementation obviously isn't a normal distribution VAE -- but maybe the author had a different goal in mind.
– Sycorax
Commented Jan 12, 2020 at 0:59
• Yes. I think that makes sense. What about batch normalization to mean and variance layer Commented Jan 12, 2020 at 3:59
• github.com/bjlkeng/sandbox/blob/master/notebooks/… this author apply batch normal to mean layer. What I am confused is the mean distribution is already similar to unit normal. Commented Jan 12, 2020 at 4:01

Batch norm as the last layer of the encoder isn't technically wrong, but it is likely to be a bad idea (in general, never use batch norm as the last layer). And you can see in the github link referenced, that the results from that model were pretty poor due to this.

Relu for std/variance could be valid, if you decided to directly predict std/variance instead of log variance.

Relu for mean is wrong.

Of course that's not to say that the image of code you attached is wrong, since there's not enough context to tell.

• So if the context is not clear, can I say 1: not to use batch norm and 2: no activation for the last layer of encoder will be a way to go? Commented Jan 13, 2020 at 17:44
• yes, that's right Commented Jan 13, 2020 at 17:52

I read this description from the document of TensorFlow:

Note, it's common practice to avoid using batch normalization when training VAEs, since the additional stochasticity due to using mini-batches may aggravate instability on top of the stochasticity from sampling.

https://www.tensorflow.org/tutorials/generative/cvae#network_architecture