I have been reading this paper: https://arxiv.org/pdf/1312.6114.pdf 117 to build a variational autoencoder. Then, I stumbled upon the VAE example that pytorch offers: https://github.com/pytorch/examples/blob/master/vae/main.py#L39 158. This one is for binary data because it uses a Bernoulli distribution in the decoder (basically the application of a sigmoid activation function to the outputs). Below there is the part of the paper where they explicitly say so:

image.png1930×600 67.8 KB

I am more interested in real-valued  data (-∞, ∞) and need the decoder of this VAE to reconstruct  a multivariate Gaussian distribution instead. In short – how to achieve this with the Pytorch’s example above?
I am a little confused and let me elaborate about it. For real-valued cases they wrote:

image.png1934×666 80.3 KB

What I did was the following to achieve Eq. (12):
diff --git a/vae/main.py b/vae/main.py
index 6286592..246dd97 100644
--- a/vae/main.py
+++ b/vae/main.py
@@ -44,6 +44,7 @@ class VAE(nn.Module):
         self.fc21 = nn.Linear(400, 20)
         self.fc22 = nn.Linear(400, 20)
         self.fc3 = nn.Linear(20, 400)
+        self.fc31 = nn.Linear(20, 400)
         self.fc4 = nn.Linear(400, 784)

     def encode(self, x):
@@ -56,8 +57,10 @@ class VAE(nn.Module):
         return mu + eps*std

     def decode(self, z):
-        h3 = F.relu(self.fc3(z))
-        return torch.sigmoid(self.fc4(h3))
+        mu = F.relu(self.fc3(z))
+        logvar = F.relu(self.fc3(m))
+        gaussian = self.reparameterize(mu, logvar)
+        return self.fc4(gaussian)

     def forward(self, x):
         mu, logvar = self.encode(x.view(-1, 784))
@@ -71,15 +74,17 @@ optimizer = optim.Adam(model.parameters(), lr=1e-3)

 # Reconstruction + KL divergence losses summed over all elements and batch
 def loss_function(recon_x, x, mu, logvar):
-    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')

     # see Appendix B from VAE paper:
     # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
     # https://arxiv.org/abs/1312.6114
     # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
+    criterion = torch.nn.MSELoss()
+    MSE = criterion(recon_x, x) * .5
+
     KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

-    return BCE + KLD
+    return MSE + KLD


 def train(epoch):

As you see, I removed the sigmoid activation function, I reconstruct the Gaussian by taking z and calling the reparameterize function, and used MSE loss function. Is this approach correct? Or how should the decoder look like?
I would appreciate any input about this.