For a Variational Autoencoder (VAE) couldn't we just sample from a standard normal distribution without the encoder at training?

Question

Normally when training a VAE, we have an encoder and a decoder. The encoder maps the input to a series of values that represent the mean and variance of a distribution (the size of this output tensor is 2 * latent_size). After that, we sample from this (using the reparametrization trick) and the decoder uses these values to generate an output. The loss tries to both minimize the distance between the generated output and the input and to move the means and stds from the encoder close to 0 and 1 (standard distr.)

Considering the fact that at "inference" / generation, we are sampling directly from a standard normal distribution, couldn't we have just trained the decoder with samples like this and not use the encoder, which ultimately only learns to map the input to the 0, 1 pairs (from my experience by the end of the training they get quite close to these values)? We will still have the latent space regularization we need for generating coherent outputs.

Also because the KL loss part uses directly the outputs from the encoder, there is no gradient to somehow affect the decoder, so the decoder only learns from the reconstruction loss. (I might be wrong here)

A quick reference implementation of what I was thinking:

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Load MNIST dataset
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
mnist = torchvision.datasets.MNIST(root='./data', download=True, transform=transform)

batch_size = 128
train_dataset = torch.utils.data.DataLoader(mnist, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=4)

class VAE(nn.Module):
    def __init__(self, input_size, hidden_size, latent_size):
        super(VAE, self).__init__()
        self.latent_size = latent_size
        self.input_size = input_size
        self.hidden_size = hidden_size
        
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, latent_size * 2)
        self.fc3 = nn.Linear(latent_size, hidden_size)
        self.fc4 = nn.Linear(hidden_size, input_size)


    def encoder(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

    def reparameterize(self, mu, log_var):
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        return eps.mul(std).add_(mu)

    def decoder(self, z):
        z = F.relu(self.fc3(z))
        z = torch.sigmoid(self.fc4(z))
        return z

    def forward(self, x, z):
        # mu, log_var = self.encoder(x).split(self.latent_size, dim=1)
        # z = self.reparameterize(mu, log_var)
        recon_x = self.decoder(z)
        return recon_x

def loss_fn(recon_x, x):
    BCE = F.binary_cross_entropy(recon_x, x, reduction='sum')
    # KLD = -torch.mean(-0.5 * torch.sum(1 + log_var - mu ** 2 - log_var.exp(), dim = 1), dim = 0)
    return BCE

# Initialize VAE model
input_size = 784
hidden_size = 512
latent_size = 32
model = VAE(input_size, hidden_size, latent_size)
model = model.to(device)

# Choose optimizer and loss function
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)

# Train loop
num_epochs = 10
for epoch in range(num_epochs):
    train_loss = 0
    for x, _ in train_dataset:
        x = x.view(-1, 784).to(device)
        z = torch.randn(x.shape[0], latent_size)
        z = z.to(device)
        recon_x = model(x, z)
        loss = loss_fn(recon_x, x)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_loss += loss.item()

    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {train_loss / len(train_dataset.dataset)}')

Any reviews or improvements are more than welcome!

Answer 1

If you want to use the decoder as a generative model, you can theoretically input random vectors and force it to generate an image (this is what the generator in a GAN does). But that would cause, in the best case, that the model will only work for specific random vectors.

If you just not use the encoder you're basically generating 2 random vectors which output two different images. What happens if you take 2 very similar images in the training set, the same image rotated by 1 degree for example, and randomly generate two random vectors that are very far apart? Also think of the inverse case where you generate 2 infinitesimally close random vectors, but force them to generate very different images ( a white plane over blue skies and a black sheep in a green meadow).

You are basically forcing the generative model to learn contradictory information and therefore hurting performance.

The clever solution for that was the encoder. It allows the model to automatically constrain the random vectors so that similar inputs generate similar random vectors, thus allowing the decoder to learn much better. It also makes the latent space continuous so that for every random vector you get a pretty valid output.

Hope it helps somewhat.