# Autoencoder/VAE Architecture for Image Generation in Microsoft’s AI-For-Beginners

> Explore the Autoencoder VAE architecture for image generation in Microsoft's AI-For-Beginners. Learn about its dual-output encoder, reparameterization trick, and symmetric decoder.

- Repository: [Microsoft/AI-For-Beginners](https://github.com/microsoft/AI-For-Beginners)
- Tags: architecture
- Published: 2026-08-23

---

**The microsoft/AI-For-Beginners repository implements a fully-connected Variational Autoencoder featuring a dual-output encoder that predicts Gaussian parameters μ and log-variance, a reparameterization trick for latent sampling, and a symmetric decoder with sigmoid activation trained with binary cross-entropy and KL-divergence losses to generate 28×28 MNIST digits.**

Lesson 09 of the computer vision curriculum (`lessons/4-ComputerVision/09-Autoencoders/`) provides complete implementations in both TensorFlow/Keras and PyTorch. These notebooks demonstrate how a dense-layer autoencoder/VAE architecture learns compressed latent representations and synthesizes new handwritten digit images through probabilistic decoding.

## Encoder Architecture

The encoder compresses the 784-dimensional flattened MNIST input (28×28 pixels) through a series of fully-connected layers. According to the source code in `AutoencodersTF.ipynb` and `AutoEncodersPyTorch.ipynb`, the architecture follows this specification:

- **Input layer**: Accepts 784 grayscale pixel values (flattened from 28×28)
- **Hidden layer 1**: 512 units with **ReLU** activation
- **Hidden layer 2**: 512 units with **ReLU** activation
- **Output split**: Two parallel dense layers producing `mu` (mean μ) and `log_var` (log-variance) vectors

The final layers do not output a single vector but rather parameterize the latent Gaussian distribution. The latent dimension is typically configured as **2** for visualization purposes or **10** for more expressive generation capacity.

```python

# PyTorch encoder implementation (from AutoEncodersPyTorch.ipynb)

def encode(self, x):
    x = torch.relu(self.fc1(x))  # fc1: Linear(784, 512)

    x = torch.relu(self.fc2(x))  # fc2: Linear(512, 512)

    return self.fc_mu(x), self.fc_logvar(x)  # Linear(512, latent_dim)

```

## Latent Space and the Reparameterization Trick

To enable backpropagation through the stochastic sampling operation, the implementation uses the **reparameterization trick**. Instead of sampling directly from the distribution defined by μ and σ, the model samples epsilon (ε) from a standard normal distribution and transforms it deterministically.

The operation follows: `z = μ + σ·ε`, where `σ = exp(0.5 * log_var)` and `ε ~ N(0, I)`.

In the PyTorch implementation (`AutoEncodersPyTorch.ipynb`), this logic resides in a dedicated method:

```python
def reparameterize(self, mu, logvar):
    std = torch.exp(0.5 * logvar)
    eps = torch.randn_like(std)
    return mu + eps * std

```

This separates the stochastic component (epsilon) from the learnable parameters (μ and log-variance), allowing gradient flow during training.

## Decoder Architecture

The decoder mirrors the encoder structure to reconstruct the original image dimensions from the sampled latent vector **z**. The architecture specifications found in both framework implementations include:

- **Input**: Sampled latent vector `z` (dimension 2 or 10)
- **Hidden layer 1**: 512 units with **ReLU** activation
- **Hidden layer 2**: 512 units with **ReLU** activation
- **Output layer**: 784 units with **sigmoid** activation, reshaped to 28×28

The sigmoid activation ensures pixel values fall within the [0, 1] range, matching the normalized MNIST input.

```python

# TensorFlow/Keras decoder (from AutoencodersTF.ipynb)

latent_inputs = keras.Input(shape=(latent_dim,))
x = keras.layers.Dense(512, activation='relu')(latent_inputs)
x = keras.layers.Dense(512, activation='relu')(x)
outputs = keras.layers.Dense(784, activation='sigmoid')(x)
decoder = keras.Model(latent_inputs, outputs, name="decoder")

```

## Loss Function and Training Configuration

The training objective combines reconstruction accuracy with distribution regularization. As implemented in the lesson notebooks, the total loss is the sum of two components:

- **Reconstruction loss**: Binary cross-entropy between the input image and the decoder output (measuring pixel-wise reconstruction error)
- **KL Divergence**: `0.5 * sum(exp(log_var) + mu^2 - 1 - log_var)`, which forces the learned latent distribution toward a standard normal distribution

Training hyperparameters used in the source code:
- **Optimizer**: Adam with learning rate **1e-3**
- **Batch size**: 128
- **Epochs**: 30–50 (sufficient for MNIST convergence)
- **Dataset**: MNIST handwritten digits normalized to [0, 1]

## Framework-Specific Implementation Differences

While the architecture remains identical, the organizational structure differs between frameworks.

**TensorFlow/Keras** (`AutoencodersTF.ipynb`): Uses the Functional API to define clear encoder and decoder sub-models, connecting them via the reparameterization layer in the training step.

**PyTorch** (`AutoEncodersPyTorch.ipynb`): Encapsulates the logic in a `VAE` class inheriting from `nn.Module` with explicit methods:
- `encode(x)` → returns mu, logvar
- `reparameterize(mu, logvar)` → returns z
- `decode(z)` → returns reconstructed image
- `forward(x)` → orchestrates the full inference pass

Visual diagrams of this architecture are provided in the repository at `lessons/4-ComputerVision/09-Autoencoders/images/autoencoder_schema.jpg` and `lessons/4-ComputerVision/09-Autoencoders/images/vae.png`.

## Summary

- The VAE uses a **symmetric dense architecture** (784→512→512→latent→512→512→784) with ReLU activations in hidden layers and sigmoid on the output.
- The encoder outputs **μ and log-variance** parameters rather than a single deterministic vector.
- The **reparameterization trick** enables gradient descent through the stochastic latent sampling operation.
- Training optimizes a composite **binary cross-entropy + KL-divergence** loss using the Adam optimizer.
- Complete implementations are available in `AutoencodersTF.ipynb` (TensorFlow) and `AutoEncodersPyTorch.ipynb` (PyTorch) under the `09-Autoencoders` lesson directory.

## Frequently Asked Questions

### What dataset does the autoencoder/VAE implementation use?

The implementation uses the **MNIST** dataset of handwritten digits. Each image is 28×28 pixels, grayscale, and flattened to a 784-dimensional vector before being fed into the encoder.

### Why is the latent dimension typically set to 2 in the examples?

A latent dimension of **2** allows for easy visualization of the learned distribution in a 2D plane, making it straightforward to plot the μ vectors and sample points for educational purposes. The architecture supports larger dimensions (e.g., 10) for higher-quality generation.

### What is the purpose of the reparameterization trick in the VAE?

The **reparameterization trick** moves the random sampling operation outside of the main computational graph. By sampling epsilon (ε) from a standard normal distribution and transforming it via `μ + σ·ε`, the model enables backpropagation to flow through the mean and variance parameters while maintaining the stochastic nature of the latent space.

### Where can I find the complete model definitions and training loops?

The complete TensorFlow implementation is located in `lessons/4-ComputerVision/09-Autoencoders/AutoencodersTF.ipynb`, while the PyTorch version resides in `lessons/4-ComputerVision/09-Autoencoders/AutoEncodersPyTorch.ipynb`. Both files contain the full VAE class definitions, loss computations, and training epochs necessary to reproduce the image generation results.