A Rigorous Mathematical Exploration of Variational Autoencoders

09 Oct, 2024

Introduction

Variational Autoencoders (VAEs) represent a fascinating fusion of deep learning and probabilistic modeling. This post aims to provide a comprehensive exploration of VAEs, balancing rigorous mathematics with practical examples and applications. We'll journey from foundational concepts to advanced theoretical insights, making this exploration accessible to both theoretically inclined beginners, regular readers and practitioners. We will certainly try our best to make you understand how to use VAE in generating an image using a toy example.

Probabilistic Foundations

Measure-Theoretic Probability

Let $(Ω, ℱ, P)$ be a probability space, where:

$Ω$ is the sample space
$ℱ$ is a $σ$ -algebra on $Ω$
$P : ℱ \to [0, 1]$ is a probability measure

Toy Example: Imagine we're modeling the process of generating handwritten digits. Here, $Ω$ could be all possible 28x28 pixel images, $ℱ$ would be sets of these images (e.g., all images of the digit "7"), and $P$ would assign probabilities to these sets.

Random Variables and Distributions

Let $X : Ω \to 𝒳$ and $Z : Ω \to 𝒵$ be random variables representing observed and latent data respectively.

The core of VAE theory lies in the following decomposition:

p (x, z) = p (x | z) p (z)

where $p (z)$ is the prior over latent variables and $p (x | z)$ is the likelihood.

Proof: This follows directly from the definition of conditional probability:

\begin{matrix} p (x, z) & = p (x | z) p (z) (by definition of conditional probability) \\ = p (z | x) p (x) (alternative factorization) \end{matrix}

Practical Example: In our handwritten digit VAE:

$x$ could be a 28x28 image of a digit
$z$ could be a 10-dimensional vector encoding factors like stroke width, tilt, etc.
$p (z)$ might be a standard normal distribution $𝒩 (0, I)$
$p (x | z)$ would be our decoder network, generating an image given a latent code

The Variational Inference Framework

The Evidence Lower Bound (ELBO)

Theorem (Evidence Lower Bound): For any distributions $q (z | x)$ and $p (x, z)$ ,

\log p (x) \geq 𝔼_{q (z | x)} [\log p (x, z) - \log q (z | x)] \equiv ELBO (x)

Proof:

\begin{matrix} \log p (x) & = \log \int p (x, z) d z \\ = \log \int p (x, z) \frac{q (z | x)}{q (z | x)} d z \\ = \log 𝔼_{q (z | x)} [\frac{p (x, z)}{q (z | x)}] \\ \geq 𝔼_{q (z | x)} [\log \frac{p (x, z)}{q (z | x)}] (by Jensen's inequality) \end{matrix}

Practical Implementation: In PyTorch, we might implement the ELBO loss as:

def elbo_loss(x, x_recon, mu, logvar):
    recon_loss = F.mse_loss(x_recon, x, reduction='sum')
    kl_div = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon_loss + kl_div

ELBO Decomposition

We can further decompose the ELBO:

ELBO (x) = 𝔼_{q (z | x)} [\log p (x | z)] - KL (q (z | x) | | p (z))

Proof:

\begin{matrix} ELBO (x) & = 𝔼_{q (z | x)} [\log p (x, z) - \log q (z | x)] \\ = 𝔼_{q (z | x)} [\log p (x | z) + \log p (z) - \log q (z | x)] \\ = 𝔼_{q (z | x)} [\log p (x | z)] + 𝔼_{q (z | x)} [\log p (z) - \log q (z | x)] \\ = 𝔼_{q (z | x)} [\log p (x | z)] - KL (q (z | x) | | p (z)) \end{matrix}

Application: In image generation, this decomposition shows us we're balancing two objectives:

Reconstruction accuracy: How well can we reconstruct the input image?
Latent space regularity: How close is our encoded distribution to the prior?

The Reparameterization Trick

Theorem (Reparameterization Trick): Let $g_{ϕ} (ϵ, x)$ be a differentiable transformation. Then,

𝔼_{q_{ϕ} (z | x)} [f (z)] = 𝔼_{ϵ ~ p (ϵ)} [f (g_{ϕ} (ϵ, x))]

where $ϵ$ is drawn from some fixed distribution $p (ϵ)$ .

Proof: Let $z = g_{ϕ} (ϵ, x)$ where $ϵ ~ p (ϵ)$ . Then,

\begin{matrix} 𝔼_{q_{ϕ} (z | x)} [f (z)] & = \int f (z) q_{ϕ} (z | x) d z \\ = \int f (g_{ϕ} (ϵ, x)) p (ϵ) d ϵ (by change of variables) \\ = 𝔼_{ϵ ~ p (ϵ)} [f (g_{ϕ} (ϵ, x))] \end{matrix}

PyTorch Implementation:

class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()

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

    def forward(self, x):
        mu, logvar = self.encoder(x)
        z = self.reparameterize(mu, logvar)
        return self.decoder(z), mu, logvar

Advanced ELBO Analysis

Theorem (Information-Theoretic ELBO Decomposition): The ELBO can be expressed as:

ELBO = 𝔼_{p_{data} (x)} [𝔼_{q_{ϕ} (z | x)} [\log p_{θ} (x | z)]] - I_{q} (x; z) + H (x)

where $I_{q} (x; z)$ is the mutual information between $x$ and $z$ under $q_{ϕ} (x, z) = p_{data} (x) q_{ϕ} (z | x)$ , and $H (x)$ is the entropy of the data distribution.

Proof:

\begin{matrix} ELBO & = 𝔼_{p_{data} (x)} [𝔼_{q_{ϕ} (z | x)} [\log p_{θ} (x | z)] - KL (q_{ϕ} (z | x) | | p (z))] \\ = 𝔼_{p_{data} (x)} [𝔼_{q_{ϕ} (z | x)} [\log p_{θ} (x | z)]] - 𝔼_{p_{data} (x)} [KL (q_{ϕ} (z | x) | | p (z))] \\ = 𝔼_{p_{data} (x)} [𝔼_{q_{ϕ} (z | x)} [\log p_{θ} (x | z)]] - (I_{q} (x; z) + 𝔼_{p_{data} (x)} [KL (q_{ϕ} (z | x) | | q_{ϕ} (z))]) \\ = 𝔼_{p_{data} (x)} [𝔼_{q_{ϕ} (z | x)} [\log p_{θ} (x | z)]] - I_{q} (x; z) + H (x) \end{matrix}

Application: This decomposition reveals insights for disentangled representation learning. By explicitly controlling $I_{q} (x; z)$ , we can encourage the model to learn more interpretable latent factors.

Normalizing Flows and VAEs

Theorem (Change of Variables Formula for Normalizing Flows): Let $z_{K} = f_{K} \circ . . . \circ f_{1} (z_{0})$ be a normalizing flow. Then,

\log q_{K} (z_{K}) = \log q_{0} (z_{0}) - \sum_{k = 1}^{K} \log | det \frac{\partial f_{k}}{\partial z_{k - 1}} |

Proof: By the change of variables formula and the chain rule:

\begin{matrix} q_{K} (z_{K}) & = q_{0} (z_{0}) \prod_{k = 1}^{K} | det \frac{\partial f_{k}^{- 1}}{\partial z_{k}} | \\ \log q_{K} (z_{K}) & = \log q_{0} (z_{0}) + \sum_{k = 1}^{K} \log | det \frac{\partial f_{k}^{- 1}}{\partial z_{k}} | \\ = \log q_{0} (z_{0}) - \sum_{k = 1}^{K} \log | det \frac{\partial f_{k}}{\partial z_{k - 1}} | \end{matrix}

Practical Example: Implementing a simple planar flow in PyTorch:

class PlanarFlow(nn.Module):
    def __init__(self, dim):
        super(PlanarFlow, self).__init__()
        self.w = nn.Parameter(torch.randn(dim))
        self.b = nn.Parameter(torch.randn(1))
        self.u = nn.Parameter(torch.randn(dim))

    def forward(self, z):
        linear_term = torch.sum(self.w * z, dim=1, keepdim=True) + self.b
        return z + self.u * torch.tanh(linear_term)

    def log_det_jacobian(self, z):
        linear_term = torch.sum(self.w * z, dim=1, keepdim=True) + self.b
        psi = (1 - torch.tanh(linear_term)**2) * self.w
        return torch.log(torch.abs(1 + torch.sum(psi * self.u, dim=1, keepdim=True)))

Tighter Variational Bounds

Theorem (IWAE Bound): For any number of samples $K$ ,

\log p (x) \geq 𝔼_{z_{1}, . . ., z_{K} ~ q_{ϕ} (z | x)} [\log \frac{1}{K} \sum_{k = 1}^{K} \frac{p_{θ} (x, z_{k})}{q_{ϕ} (z_{k} | x)}] \equiv ℒ_{K}

Proof: By Jensen's inequality and the concavity of log:

\begin{matrix} \log p (x) & = \log 𝔼_{z ~ q_{ϕ} (z | x)} [\frac{p_{θ} (x, z)}{q_{ϕ} (z | x)}] \\ = \log 𝔼_{z_{1}, . . ., z_{K} ~ q_{ϕ} (z | x)} [\frac{1}{K} \sum_{k = 1}^{K} \frac{p_{θ} (x, z_{k})}{q_{ϕ} (z_{k} | x)}] \\ \geq 𝔼_{z_{1}, . . ., z_{K} ~ q_{ϕ} (z | x)} [\log \frac{1}{K} \sum_{k = 1}^{K} \frac{p_{θ} (x, z_{k})}{q_{ϕ} (z_{k} | x)}] \end{matrix}

PyTorch Implementation:

def iwae_loss(x, model, K=10):
    x_expanded = x.unsqueeze(0).expand(K, *x.shape)
    x_recon, mu, logvar = model(x_expanded)
    
    z = model.reparameterize(mu, logvar)
    log_p_x_given_z = -F.mse_loss(x_recon, x_expanded, reduction='none').sum(dim=(1,2,3))
    log_p_z = -0.5 * z.pow(2).sum(dim=1)
    log_q_z_given_x = -0.5 * (logvar + (z - mu).pow(2) / logvar.exp()).sum(dim=1)
    
    log_w = log_p_x_given_z + log_p_z - log_q_z_given_x
    return -torch.logsumexp(log_w, dim=0) + np.log(K)

Mathematical Model of a VAE for Image Generation

Let's develop a rigorous mathematical model for a VAE designed to generate images, accompanied by a toy example to illustrate the concepts.

Data and Latent Space

Let $𝒳 = [0, 1]^{n^{2}}$ be our data space, representing all possible $n \times n$ grayscale images. Each $x \in 𝒳$ is a flattened vector of pixel intensities.

Let $𝒵 = ℝ^{d}$ be our latent space, where $d$ is the dimensionality of the latent representation.

Toy Example: Let's consider a very simple case where:

Our images are 3x3 grayscale (n = 3)
Our latent space is 2-dimensional (d = 2)

So, $𝒳 = [0, 1]^{9}$ and $𝒵 = ℝ^{2}$

A sample image might look like:

[0.1, 0.5, 0.9]
[0.2, 0.6, 0.8]
[0.3, 0.7, 0.4]

Which we'd flatten to: x = [0.1, 0.5, 0.9, 0.2, 0.6, 0.8, 0.3, 0.7, 0.4]

Generative Model

We define our generative model as follows:

Prior: $p (z) = 𝒩 (0, I_{d})$ , where $I_{d}$ is the $d$ -dimensional identity matrix.
Decoder: $p_{θ} (x | z) = 𝒩 (μ_{θ} (z), σ^{2} I_{n^{2}})$ , where:
- $μ_{θ} : ℝ^{d} \to [0, 1]^{n^{2}}$ is a neural network with parameters $θ$
- $σ^{2}$ is a fixed variance for simplicity (though this could also be learned)

Toy Example: Let's define a simple decoder network:

$μ_{θ} (z) = sigmoid (W_{2} \cdot ReLU (W_{1} z + b_{1}) + b_{2})$

Where: $W_{1} \in ℝ^{4 \times 2}, b_{1} \in ℝ^{4}, W_{2} \in ℝ^{9 \times 4}, b_{2} \in ℝ^{9}$

Let $σ^{2} = 0.1$

Inference Model

Our inference model (encoder) is defined as:

$q_{ϕ} (z | x) = 𝒩 (μ_{ϕ} (x), Σ_{ϕ} (x))$

where:

$μ_{ϕ} : [0, 1]^{n^{2}} \to ℝ^{d}$ is a neural network with parameters $ϕ$
$Σ_{ϕ} : [0, 1]^{n^{2}} \to ℝ^{d \times d}$ is a neural network outputting a diagonal covariance matrix

Toy Example: Let's define a simple encoder network:

$[μ_{ϕ} (x), \log Σ_{ϕ} (x)] = W_{4} \cdot ReLU (W_{3} x + b_{3}) + b_{4}$

Where: $W_{3} \in ℝ^{6 \times 9}, b_{3} \in ℝ^{6}, W_{4} \in ℝ^{4 \times 6}, b_{4} \in ℝ^{4}$

The output is split into $μ_{ϕ} (x) \in ℝ^{2}$ and $\log Σ_{ϕ} (x) \in ℝ^{2}$ (diagonal elements of the covariance matrix in log space).

7.4 ELBO for Image Generation

The Evidence Lower Bound (ELBO) for our image generation VAE is:

ELBO (x) = 𝔼_{q_{ϕ} (z | x)} [\log p_{θ} (x | z)] - KL (q_{ϕ} (z | x) | | p (z))

Let's expand each term:

Reconstruction term:

𝔼_{q_{ϕ} (z | x)} [\log p_{θ} (x | z)] & = 𝔼_{q_{ϕ} (z | x)} [- \frac{n^{2}}{2} \log (2 π σ^{2}) - \frac{1}{2 σ^{2}} | | x - μ_{θ} (z) | |^{2}] & \approx - \frac{n^{2}}{2} \log (2 π σ^{2}) - \frac{1}{2 σ^{2}} | | x - μ_{θ} (μ_{ϕ} (x)) | |^{2}

where we've used a single sample approximation for the expectation.

KL divergence term:

KL (q_{ϕ} (z | x) | | p (z)) = \frac{1}{2} (tr (Σ_{ϕ} (x)) + μ_{ϕ} (x)^{T} μ_{ϕ} (x) - d - \log det (Σ_{ϕ} (x)))

Toy Example Calculation: Let's calculate the ELBO for our sample image x = [0.1, 0.5, 0.9, 0.2, 0.6, 0.8, 0.3, 0.7, 0.4]

Encode: Suppose $μ_{ϕ} (x) = [0.5, - 0.3]$ and $\log Σ_{ϕ} (x) = [- 1.2, - 0.8]$
Sample z: $z = [0.4, - 0.5]$ (sampled from $𝒩 (μ_{ϕ} (x), Σ_{ϕ} (x))$ )
Decode: Suppose $μ_{θ} (z) = [0.15, 0.45, 0.85, 0.25, 0.55, 0.75, 0.35, 0.65, 0.45]$
Reconstruction term: $- \frac{9}{2} \log (2 π \cdot 0.1) - \frac{1}{2 \cdot 0.1} (0.1 - 0.15)^{2} + . . . + (0.4 - 0.45)^{2} \approx - 10.5$
KL divergence term: $\frac{1}{2} (e^{- 1.2} + e^{- 0.8} + {0.5}^{2} + (- 0.3)^{2} - 2 + 1.2 + 0.8) \approx 0.3$
ELBO: $ELBO (x) \approx - 10.5 - 0.3 = - 10.8$

Optimization Objective

Our final optimization objective is to maximize the ELBO:

{max}_{θ, ϕ} 𝔼_{p_{data} (x)} [ELBO (x)]

which is equivalent to minimizing the negative ELBO:

{min}_{θ, ϕ} 𝔼_{p_{data} (x)} [\frac{1}{2 σ^{2}} | | x - μ_{θ} (μ_{ϕ} (x)) | |^{2} + \frac{1}{2} (tr (Σ_{ϕ} (x)) + μ_{ϕ} (x)^{T} μ_{ϕ} (x) - d - \log det (Σ_{ϕ} (x)))]

Toy Example: In practice, we would use stochastic gradient descent to optimize this objective over a dataset of images. For our single image example, one step of gradient descent might look like:

θ \leftarrow θ + α \nabla_{θ} ELBO (x)

ϕ \leftarrow ϕ + α \nabla_{ϕ} ELBO (x)

Where $α$ is the learning rate.

Image Generation Process

To generate a new image:

Sample $z ~ 𝒩 (0, I_{d})$
Compute $μ_{θ} (z)$
Sample $x ~ 𝒩 (μ_{θ} (z), σ^{2} I_{n^{2}})$
Reshape the resulting vector into an $n \times n$ image

Toy Example:

Sample $z = [0.1, 0.7]$ from $𝒩 (0, I_{2})$
Compute $μ_{θ} (z) = [0.2, 0.4, 0.6, 0.3, 0.5, 0.7, 0.1, 0.8, 0.9]$
Sample $x$ from $𝒩 (μ_{θ} (z), 0.1 I_{9})$ , say we get: $x = [0.25, 0.35, 0.55, 0.35, 0.45, 0.75, 0.15, 0.85, 0.95]$

Reshape into a 3x3 image:

[0.25, 0.35, 0.55]
[0.35, 0.45, 0.75]
[0.15, 0.85, 0.95]

Theoretical Properties

Consistency: As the number of training samples approaches infinity and with sufficient model capacity, the VAE converges to the true data distribution.
Manifold Learning: The VAE learns a low-dimensional manifold in $𝒵$ that captures the structure of the image data.
Disentanglement: Under certain conditions (e.g., using a β-VAE formulation), the latent space can capture disentangled factors of variation in the images.
Sample Quality vs. Reconstruction Quality Trade-off: There's an inherent tension between generating high-quality samples and accurately reconstructing inputs, controlled by the weight of the KL divergence term.

Toy Example Illustration: In our 2D latent space, we might find that:

One dimension controls the overall brightness of the image
The other dimension controls the contrast between the center and edges

This would be an example of learning a disentangled representation, where each latent dimension corresponds to a meaningful factor of variation in the data.

Conclusion

Through this detailed exploration of VAEs for image generation, complete with a toy example, we've seen how these models bridge the gap between deep learning and probabilistic modeling. The mathematical formulation provides a rigorous foundation, while the toy example offers concrete insights into how VAEs operate in practice.

This combination of theory and example illustrates the power of VAEs in learning compact, meaningful representations of complex data like images. As we continue to advance both the theory and application of VAEs, we open up new possibilities in generative modeling, representation learning, and beyond.