E2E: Variational Autoencoder - End-to-End Implementations

Implement E2E: Variational Autoencoder in Python with NumPy. hard exercise with examples, hints, and instant feedback. No setup or signup required.

End-to-End Variational Autoencoder

What you'll build

A complete VAE in NumPy: an MLP encoder that outputs a Gaussian per input, the reparameterization trick, an MLP decoder, the negative-ELBO loss, and sampling from the prior. It assembles VAE Reparameterization Trick and VAE Loss Function from Generative Models into one model.

Architecture

Stage	Operation	Shape
input	`x`	`(batch, input_dim)`
encoder hidden	`h = relu(x @ W_enc + b_enc)`	`(batch, hidden_dim)`
latent mean	`mu = h @ W_mu + b_mu`	`(batch, latent_dim)`
latent log-variance	`log_var = h @ W_logvar + b_logvar`	`(batch, latent_dim)`
sample	`z = mu + exp(0.5·log_var) ⊙ ε`, `ε = np.random.randn(*mu.shape)`	`(batch, latent_dim)`
decoder hidden	`h' = relu(z @ W_dec1 + b_dec1)`	`(batch, hidden_dim)`
reconstruction	`x_recon = sigmoid(h' @ W_dec2 + b_dec2)`	`(batch, input_dim)`, values in (0, 1)

The constructor (given) seeds np.random.seed(42) and initializes every weight as N(0, 0.01²) with zero biases.

Loss (negative ELBO)

The same convention as VAE Loss Function (the standard one from Kingma & Welling and the PyTorch VAE example):

recon_loss = mean((x − x_recon)²)                                        # over all elements
kl_loss    = mean_i [ −0.5 · Σ_j (1 + log_var_ij − mu_ij² − exp(log_var_ij)) ]   # sum over latent dims, mean over batch
total_loss = recon_loss + kl_loss

Earlier versions of this problem averaged the KL over latent dims as well; that shrinks it by a factor of latent_dim relative to the reconstruction term. The convention above is the one both problems now use.

Method contract

Method	Returns
`encode(x)`	tuple `(mu, log_var)`, each `(batch, latent_dim)`
`reparameterize(mu, log_var)`	`z`, drawing ε once with `np.random.randn(*mu.shape)` (no re-seeding)
`decode(z)`	`x_recon` `(batch, input_dim)`, sigmoid output
`forward(x)`	dict `{'mu', 'log_var', 'z', 'x_recon'}`
`compute_loss(x, x_recon, mu, log_var)`	dict `{'total_loss', 'recon_loss', 'kl_loss'}` of floats
`sample(num_samples)`	`decode(z)` with `z = np.random.randn(num_samples, latent_dim)`: `(num_samples, input_dim)`

Edge cases and conventions

With the tiny initial weights, mu ≈ 0 and log_var ≈ 0, so the initial KL is close to 0 and x_recon ≈ 0.5.
kl_loss ≥ 0 always; total_loss is exactly the sum of the other two.

Constraints

NumPy only, vectorised over the batch. Training (backprop) is not required here; the PyTorch track's autograd makes it a one-liner.

Follow-up

With sigmoid outputs and binary data, why is BCE the more principled reconstruction loss?
How do you detect posterior collapse (KL → 0 for every latent dim), and what do KL annealing and free bits do?

Examples

Input: vae.compute_loss(x=[[1.0, 2.0]], x_recon=[[0.0, 0.0]], mu=[[0.5, -0.5]], log_var=[[0.0, 1.0]])
Output: {'total_loss': 3.1091, 'recon_loss': 2.5, 'kl_loss': 0.6091} (rounded)

Same numbers as the VAE Loss Function example: recon = 5/2, KL summed over the 2 latent dims

Input: VAE(784, 256, 20).forward(np.random.randn(4, 784))
Output: {'mu': (4, 20), 'log_var': (4, 20), 'z': (4, 20), 'x_recon': (4, 784)}

Shapes of the returned arrays

Practice in Python

import numpy as np

class VAE:
    def __init__(self, input_dim: int, hidden_dim: int, latent_dim: int):
        """
        Initialize VAE.

        Args:
            input_dim: Input dimension (e.g., 784 for MNIST)
            hidden_dim: Hidden layer dimension
            latent_dim: Latent space dimension
        """
        np.random.seed(42)

        # Encoder weights
        self.W_enc = np.random.randn(input_dim, hidden_dim) * 0.01
        self.b_enc = np.zeros(hidden_dim)

        # Latent space (mean and log variance)
        self.W_mu = np.random.randn(hidden_dim, latent_dim) * 0.01
        self.b_mu = np.zeros(latent_dim)
        self.W_logvar = np.random.randn(hidden_dim, latent_dim) * 0.01
        self.b_logvar = np.zeros(latent_dim)

        # Decoder weights
        self.W_dec1 = np.random.randn(latent_dim, hidden_dim) * 0.01
        self.b_dec1 = np.zeros(hidden_dim)
        self.W_dec2 = np.random.randn(hidden_dim, input_dim) * 0.01
        self.b_dec2 = np.zeros(input_dim)

        self.latent_dim = latent_dim

    def encode(self, x: np.ndarray) -> tuple:
        """
        Encode input to latent distribution parameters.
        h = relu(x @ W_enc + b_enc); mu = h @ W_mu + b_mu; log_var = h @ W_logvar + b_logvar

        Returns:
            (mu, log_var): each (batch, latent_dim)
        """
        # Your code here
        pass

    def reparameterize(self, mu: np.ndarray, log_var: np.ndarray) -> np.ndarray:
        """
        Reparameterization trick: z = mu + exp(0.5 * log_var) * eps,
        with eps = np.random.randn(*mu.shape) (drawn once, no re-seeding).
        """
        # Your code here
        pass

    def decode(self, z: np.ndarray) -> np.ndarray:
        """
        Decode latent vectors: relu(z @ W_dec1 + b_dec1), then sigmoid(... @ W_dec2 + b_dec2).

        Returns:
            x_recon: (batch, input_dim), values in (0, 1)
        """
        # Your code here
        pass

    def forward(self, x: np.ndarray) -> dict:
        """
        Full forward pass.

        Returns:
            {'mu', 'log_var', 'z', 'x_recon'}
        """
        # Your code here
        pass

    def compute_loss(self, x: np.ndarray, x_recon: np.ndarray,
                     mu: np.ndarray, log_var: np.ndarray) -> dict:
        """
        Negative ELBO.
        recon_loss = MSE over all elements;
        kl_loss = -0.5 * sum over latent dims of (1 + log_var - mu^2 - exp(log_var)), averaged over the batch.

        Returns:
            {'total_loss', 'recon_loss', 'kl_loss'} as floats, total = recon + kl
        """
        # Your code here
        pass

    def sample(self, num_samples: int) -> np.ndarray:
        """
        Decode z = np.random.randn(num_samples, latent_dim) drawn from the prior.

        Returns:
            (num_samples, input_dim)
        """
        # Your code here
        pass

The interactive editor loads with Python and NumPy. Run your implementation against 8 tests, or use the progressive hints.

More exercises in this section →