tom:

I think you want
h = F.relu(self.fc3(z)) # KW use tanh instead of relu
mu = self.fc4(h)
sigma = self.fc5(h)
return mu, sigma

and then pass mu, sigma and x into your criterion to give the negative log likelihood of x as a vector sampled from N(mu_i, sigma_i^2) (i.e. you need to take the log PDF).


I hope I have understood correctly. I changed the VAE and now looks this way:
diff --git a/vae/main.py b/vae/main.py
index 6286592..606977b 100644
--- a/vae/main.py
+++ b/vae/main.py
@@ -45,6 +45,7 @@ class VAE(nn.Module):
         self.fc22 = nn.Linear(400, 20)
         self.fc3 = nn.Linear(20, 400)
         self.fc4 = nn.Linear(400, 784)
+        self.fc5 = nn.Linear(400, 784)
 
     def encode(self, x):
         h1 = F.relu(self.fc1(x))
@@ -56,13 +57,17 @@ class VAE(nn.Module):
         return mu + eps*std
 
     def decode(self, z):
-        h3 = F.relu(self.fc3(z))
-        return torch.sigmoid(self.fc4(h3))
+        h = F.relu(self.fc3(z))
+        mu = self.fc4(h)
+        logvar = self.fc5(h)
+        recon_x = self.reparameterize(mu, logvar)
+        return recon_x, mu, logvar
 
     def forward(self, x):
         mu, logvar = self.encode(x.view(-1, 784))
         z = self.reparameterize(mu, logvar)
-        return self.decode(z), mu, logvar
+        recon_x, mu, logvar = self.decode(z)
+        return recon_x, mu, logvar
 
 
 model = VAE().to(device)
@@ -71,7 +76,8 @@ 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')
+    criterion = torch.nn.MSELoss()
+    MSE = criterion(recon_x, x.view(-1, 784)) * .5
 
     # see Appendix B from VAE paper:
     # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
@@ -79,7 +85,7 @@ def loss_function(recon_x, x, mu, logvar):
     # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
     KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
 
-    return BCE + KLD
+    return MSE + KLD
 
 
 def train(epoch):
@@ -127,6 +133,8 @@ if __name__ == "__main__":
         test(epoch)
         with torch.no_grad():
             sample = torch.randn(64, 20).to(device)
-            sample = model.decode(sample).cpu()
+            sample, mu, logvar = model.decode(sample)
+            sample = sample.cpu()
+


The mu and logvar used in the loss function now come from the decoder, and in order to reconstruct X, I use self.reparameterize (not sure about this). Is this correct? I am sorry if I have misunderstood. At least these changes reduce the loss see the Bernoulli and Gaussian outputs here -> https://gist.github.com/muammar/0c0c0c53f351c85c0680017a8c41ce62 22.


So the Gaussian at the reconstruction step has nothing to do (well, except being conditional on the latents) with the Gaussian from the latents (which is the bit where you do the reparametrization and things).
I think what might cause the confusion that the “reconstruction error” is easily taken too literally when looking at the example implementation. In the paper they only take and maximize the (log) likelihood of the input data under the output distribution rather than doing actual reconstruction (by sampling or somesuch). In evaluation mode, you would be expected to sample from the Gaussian.
The last paragraph of section 2.3 in the paper talks about how in the probabilistic VAE framework this term is a […] reconstruction error in auto-encoder parlance (notice the indefinite article) or maybe rather the analogue of the reconstruction error.


Thanks for this clarification @tom.