Math Foundations: SVD, Eigendecomposition, Matrix Calculus, and ELBO Derivation

4735 words
24 min read
Python Week 1: the first filter for runtime behavior
Visual companion
Python
Type and operator map

Python Week 1: the first filter for runtime behavior

View
Revision summary

What this note is really saying

Short form

# Math Foundations: SVD, Eigendecomposition, Matrix Calculus, and ELBO Derivation ## ๐ŸŽฏ Learning Objectives - Compute and interpret singular value decomposition (SVD) and eigendecomposition - Understand matrix calculus for gradient-based optimization - Derive the ELBO from first principles with full mathematical det...

Math Foundations: SVD, Eigendecomposition, Matrix Calculus, and ELBO Derivation

๐ŸŽฏ Learning Objectives

  • Compute and interpret singular value decomposition (SVD) and eigendecomposition
  • Understand matrix calculus for gradient-based optimization
  • Derive the ELBO from first principles with full mathematical detail
  • Apply optimization concepts to generative model training
  • Connect linear algebra to key generative model operations

๐Ÿ“‹ Prerequisites

  • Linear algebra basics: Vectors, matrices, matrix multiplication
  • Calculus: Partial derivatives, chain rule
  • Probability: Expectation, Bayes rule, KL divergence
  • Generative Models Overview (Week 1): Basic ML concepts

1. ๐Ÿ“– Core Content

1.1 Intuition: Why Math Matters for Generative AI

Every generative model is built on mathematical foundations:
  • Linear algebra powers the transformations in neural networks โ€” matrix multiplications, convolutions, attention mechanisms. SVD and eigendecomposition are used for PCA (preprocessing), understanding latent spaces, and analyzing optimization landscapes.
  • Matrix calculus enables gradient-based learning โ€” backpropagation is repeated application of the chain rule on matrix expressions.
  • Optimization theory explains why certain training strategies work โ€” SGD, Adam, learning rate schedules all exploit properties of the loss landscape.
  • The ELBO is the central objective for VAEs, diffusion models, and many other generative models. Understanding its derivation is essential for debugging and extending these models.

1.2 Eigenvalues and Eigenvectors

1.2.1 Definition

For a square matrix AโˆˆRnร—nA \in \mathbb{R}^{n \times n}, a non-zero vector vโˆˆRnv \in \mathbb{R}^n is an eigenvector with corresponding eigenvalue ฮป\lambda if:
Av=ฮปvA v = \lambda v
This means: multiplying AA with vv is equivalent to scaling vv by ฮป\lambda. The eigenvector's direction is preserved; only its magnitude changes.

1.2.2 Intuition

Think of AA as a transformation that stretches/compresses space. Eigenvectors are the "special directions" that are only scaled, not rotated. The eigenvalue tells you how much stretching occurs along that direction.
  • โˆฃฮปโˆฃ>1|\lambda| > 1: Stretching
  • โˆฃฮปโˆฃ<1|\lambda| < 1: Compression
  • โˆฃฮปโˆฃ=1|\lambda| = 1: Unchanged
  • ฮป=0\lambda = 0: Direction is in the nullspace

1.2.3 Eigendecomposition

If AA has nn linearly independent eigenvectors, we can diagonalize it:
A=Vฮ›Vโˆ’1A = V \Lambda V^{-1}
where:
  • VV: Matrix whose columns are eigenvectors
  • ฮ›\Lambda: Diagonal matrix of eigenvalues diag(ฮป1,...,ฮปn)\text{diag}(\lambda_1, ..., \lambda_n) For real symmetric matrices (A=ATA = A^T), eigenvectors are orthogonal, so Vโˆ’1=VTV^{-1} = V^T:
A=Vฮ›VTA = V \Lambda V^T

Worked Example 1: Eigendecomposition of a 2ร—2 Matrix

Find eigenvalues and eigenvectors of
A=[4113]A = \begin{bmatrix} 4 & 1 \\ 1 & 3 \end{bmatrix}
. Step 1: Solve characteristic equation detโก(Aโˆ’ฮปI)=0\det(A - \lambda I) = 0:
detโก([4โˆ’ฮป113โˆ’ฮป])=(4โˆ’ฮป)(3โˆ’ฮป)โˆ’1=0\det\left(\begin{bmatrix} 4-\lambda & 1 \\ 1 & 3-\lambda \end{bmatrix}\right) = (4-\lambda)(3-\lambda) - 1 = 0 12โˆ’7ฮป+ฮป2โˆ’1=012 - 7\lambda + \lambda^2 - 1 = 0 ฮป2โˆ’7ฮป+11=0\lambda^2 - 7\lambda + 11 = 0 ฮป=7ยฑ49โˆ’442=7ยฑ52\lambda = \frac{7 \pm \sqrt{49 - 44}}{2} = \frac{7 \pm \sqrt{5}}{2} ฮป1=7+52โ‰ˆ4.618,ฮป2=7โˆ’52โ‰ˆ2.382\lambda_1 = \frac{7 + \sqrt{5}}{2} \approx 4.618, \quad \lambda_2 = \frac{7 - \sqrt{5}}{2} \approx 2.382
Step 2: Find eigenvector for ฮป1\lambda_1: (Aโˆ’ฮป1I)v1=0(A - \lambda_1 I)v_1 = 0
[4โˆ’4.618113โˆ’4.618]v1=[โˆ’0.61811โˆ’1.618]v1=0\begin{bmatrix} 4-4.618 & 1 \\ 1 & 3-4.618 \end{bmatrix} v_1 = \begin{bmatrix} -0.618 & 1 \\ 1 & -1.618 \end{bmatrix} v_1 = 0
From first row: โˆ’0.618v11+v12=0โ€…โ€ŠโŸนโ€…โ€Šv12=0.618v11-0.618v_{11} + v_{12} = 0 \implies v_{12} = 0.618v_{11} So
v1=[10.618]v_1 = \begin{bmatrix} 1 \\ 0.618 \end{bmatrix}
, normalized:
v1=11.382[10.618]โ‰ˆ[0.8510.526]v_1 = \frac{1}{\sqrt{1.382}} \begin{bmatrix} 1 \\ 0.618 \end{bmatrix} \approx \begin{bmatrix} 0.851 \\ 0.526 \end{bmatrix}
Step 3: Similarly for ฮป2\lambda_2:
v2=[โˆ’0.5260.851]v_2 = \begin{bmatrix} -0.526 \\ 0.851 \end{bmatrix}
Step 4: Verify: VTAV=ฮ›V^T A V = \Lambda
V=[0.851โˆ’0.5260.5260.851],VTAV=[4.618002.382]=ฮ›V = \begin{bmatrix} 0.851 & -0.526 \\ 0.526 & 0.851 \end{bmatrix}, \quad V^T A V = \begin{bmatrix} 4.618 & 0 \\ 0 & 2.382 \end{bmatrix} = \Lambda

1.3 Singular Value Decomposition (SVD)

SVD generalizes eigendecomposition to non-square matrices. Any matrix AโˆˆRmร—nA \in \mathbb{R}^{m \times n} can be decomposed as:
A=UฮฃVTA = U \Sigma V^T
where:
  • UโˆˆRmร—mU \in \mathbb{R}^{m \times m}: Left singular vectors (orthogonal columns, UTU=IU^T U = I)
  • ฮฃโˆˆRmร—n\Sigma \in \mathbb{R}^{m \times n}: Diagonal matrix of singular values ฯƒ1โ‰ฅฯƒ2โ‰ฅ...โ‰ฅฯƒr>0\sigma_1 \geq \sigma_2 \geq ... \geq \sigma_r > 0
  • VโˆˆRnร—nV \in \mathbb{R}^{n \times n}: Right singular vectors (orthogonal columns, VTV=IV^T V = I)
  • rr: Rank of AA (number of non-zero singular values)

1.3.1 Intuition

SVD reveals what AA does as a linear transformation:
  1. VTV^T rotates the input to align with a new coordinate system
  2. ฮฃ\Sigma scales each coordinate by ฯƒi\sigma_i
  3. UU rotates the result to produce the output (Diagram)

1.3.2 Relationship to Eigendecomposition

For ATAA^T A (a square symmetric matrix):
ATA=(UฮฃVT)T(UฮฃVT)=VฮฃTฮฃVT=Vฮ›VTA^T A = (U \Sigma V^T)^T (U \Sigma V^T) = V \Sigma^T \Sigma V^T = V \Lambda V^T
So:
  • VV contains eigenvectors of ATAA^T A
  • ฯƒi2=ฮปi\sigma_i^2 = \lambda_i (singular values squared = eigenvalues of ATAA^T A) Similarly, UU contains eigenvectors of AATA A^T.

Worked Example 2: SVD of a 2ร—2 Matrix

Compute SVD of
A=[3113]A = \begin{bmatrix} 3 & 1 \\ 1 & 3 \end{bmatrix}
. Step 1: Compute
ATA=A2=[106610]A^T A = A^2 = \begin{bmatrix} 10 & 6 \\ 6 & 10 \end{bmatrix}
Step 2: Find eigenvalues of ATAA^T A:
detโก(ATAโˆ’ฮปI)=(10โˆ’ฮป)2โˆ’36=0\det(A^T A - \lambda I) = (10-\lambda)^2 - 36 = 0 ฮป2โˆ’20ฮป+64=0\lambda^2 - 20\lambda + 64 = 0 ฮป=20ยฑ400โˆ’2562=20ยฑ122\lambda = \frac{20 \pm \sqrt{400 - 256}}{2} = \frac{20 \pm 12}{2} ฮป1=16,ฮป2=4\lambda_1 = 16, \quad \lambda_2 = 4
Singular values: ฯƒ1=16=4\sigma_1 = \sqrt{16} = 4, ฯƒ2=4=2\sigma_2 = \sqrt{4} = 2 Step 3: Find eigenvectors of ATAA^T A for VV: For ฮป1=16\lambda_1 = 16:
(ATAโˆ’16I)v=[โˆ’666โˆ’6]v=0โ€…โ€ŠโŸนโ€…โ€Šv1=12[11](A^T A - 16I)v = \begin{bmatrix} -6 & 6 \\ 6 & -6 \end{bmatrix}v = 0 \implies v_1 = \frac{1}{\sqrt{2}}\begin{bmatrix} 1 \\ 1 \end{bmatrix}
For ฮป2=4\lambda_2 = 4:
(ATAโˆ’4I)v=[6666]v=0โ€…โ€ŠโŸนโ€…โ€Šv2=12[1โˆ’1](A^T A - 4I)v = \begin{bmatrix} 6 & 6 \\ 6 & 6 \end{bmatrix}v = 0 \implies v_2 = \frac{1}{\sqrt{2}}\begin{bmatrix} 1 \\ -1 \end{bmatrix}
Step 4: Compute UU: ui=Aviฯƒiu_i = \frac{A v_i}{\sigma_i}
u1=14[3113]12[11]=142[44]=12[11]u_1 = \frac{1}{4} \begin{bmatrix} 3 & 1 \\ 1 & 3 \end{bmatrix} \frac{1}{\sqrt{2}}\begin{bmatrix} 1 \\ 1 \end{bmatrix} = \frac{1}{4\sqrt{2}}\begin{bmatrix} 4 \\ 4 \end{bmatrix} = \frac{1}{\sqrt{2}}\begin{bmatrix} 1 \\ 1 \end{bmatrix} u2=12[3113]12[1โˆ’1]=122[2โˆ’2]=12[1โˆ’1]u_2 = \frac{1}{2} \begin{bmatrix} 3 & 1 \\ 1 & 3 \end{bmatrix} \frac{1}{\sqrt{2}}\begin{bmatrix} 1 \\ -1 \end{bmatrix} = \frac{1}{2\sqrt{2}}\begin{bmatrix} 2 \\ -2 \end{bmatrix} = \frac{1}{\sqrt{2}}\begin{bmatrix} 1 \\ -1 \end{bmatrix}
Step 5: Verify:
A=UฮฃVT=12[111โˆ’1][4002]12[111โˆ’1]A = U \Sigma V^T = \frac{1}{\sqrt{2}}\begin{bmatrix} 1 & 1 \\ 1 & -1 \end{bmatrix} \begin{bmatrix} 4 & 0 \\ 0 & 2 \end{bmatrix} \frac{1}{\sqrt{2}}\begin{bmatrix} 1 & 1 \\ 1 & -1 \end{bmatrix} =[222โˆ’2]12โ‹…12[111โˆ’1]=[222โˆ’2]โ‹…12[111โˆ’1]=[3113]=A= \begin{bmatrix} 2 & 2 \\ 2 & -2 \end{bmatrix} \frac{1}{\sqrt{2}} \cdot \frac{1}{\sqrt{2}}\begin{bmatrix} 1 & 1 \\ 1 & -1 \end{bmatrix} = \begin{bmatrix} 2 & 2 \\ 2 & -2 \end{bmatrix} \cdot \frac{1}{2} \begin{bmatrix} 1 & 1 \\ 1 & -1 \end{bmatrix} = \begin{bmatrix} 3 & 1 \\ 1 & 3 \end{bmatrix} = A

1.4 Matrix Calculus

For training neural networks, we need gradients of scalar loss functions with respect to matrix parameters.

1.4.1 Key Rules

OperationGradient
โˆ‚โˆ‚x(aTx)=a\frac{\partial}{\partial x} (a^T x) = aLinear form
โˆ‚โˆ‚x(xTAx)=(A+AT)x\frac{\partial}{\partial x} (x^T A x) = (A + A^T)xQuadratic form
โˆ‚โˆ‚XTr(AX)=AT\frac{\partial}{\partial X} \text{Tr}(AX) = A^TTrace
โˆ‚โˆ‚Xdetโก(X)=detโก(X)(Xโˆ’1)T\frac{\partial}{\partial X} \det(X) = \det(X) (X^{-1})^TDeterminant
$\frac{\partial}{\partial X} \X\

Worked Example 3: Gradient of a Quadratic Form

Compute โˆ‡xf(x)\nabla_x f(x) where f(x)=(xโˆ’ฮผ)Tฮฃโˆ’1(xโˆ’ฮผ)f(x) = (x - \mu)^T \Sigma^{-1} (x - \mu) (Mahalanobis distance squared). Step 1: Expand: Let y=xโˆ’ฮผy = x - \mu, then f=yTฮฃโˆ’1yf = y^T \Sigma^{-1} y. Step 2: Using โˆ‚โˆ‚y(yTAy)=(A+AT)y\frac{\partial}{\partial y} (y^T A y) = (A + A^T)y with A=ฮฃโˆ’1A = \Sigma^{-1} (which is symmetric):
โˆ‚fโˆ‚y=2ฮฃโˆ’1y\frac{\partial f}{\partial y} = 2 \Sigma^{-1} y
Step 3: Chain rule:
โˆ‡xf=โˆ‚fโˆ‚yโ‹…โˆ‚yโˆ‚x=2ฮฃโˆ’1(xโˆ’ฮผ)โ‹…1=2ฮฃโˆ’1(xโˆ’ฮผ)\nabla_x f = \frac{\partial f}{\partial y} \cdot \frac{\partial y}{\partial x} = 2 \Sigma^{-1} (x - \mu) \cdot 1 = 2 \Sigma^{-1} (x - \mu)

1.4.2 Gradients of Log-Likelihood for Gaussians

For a multivariate Gaussian p(x)=N(xโˆฃฮผ,ฮฃ)p(x) = \mathcal{N}(x|\mu, \Sigma):
logโกp(x)=โˆ’12(xโˆ’ฮผ)Tฮฃโˆ’1(xโˆ’ฮผ)โˆ’12logโกโˆฃฮฃโˆฃโˆ’D2logโก(2ฯ€)\log p(x) = -\frac{1}{2}(x-\mu)^T \Sigma^{-1} (x-\mu) - \frac{1}{2}\log|\Sigma| - \frac{D}{2}\log(2\pi)
Gradient w.r.t. ฮผ\mu:
โˆ‡ฮผlogโกp(x)=ฮฃโˆ’1(xโˆ’ฮผ)\nabla_\mu \log p(x) = \Sigma^{-1}(x - \mu)
Gradient w.r.t. ฮฃโˆ’1\Sigma^{-1} (natural parameterization):
โˆ‡ฮฃโˆ’1logโกp(x)=12((xโˆ’ฮผ)(xโˆ’ฮผ)Tโˆ’ฮฃ)\nabla_{\Sigma^{-1}} \log p(x) = \frac{1}{2}\left((x-\mu)(x-\mu)^T - \Sigma\right)
These gradients are used in:
  • MLE for Gaussian models
  • VAE training (Gaussian encoder/decoder)
  • KL divergence computation in diffusion models

1.5 Optimization for Generative Models

1.5.1 Convex vs Non-Convex Optimization

Convex functions: Bowl-shaped, any local minimum is global. Guaranteed convergence for gradient descent.
f(ฮปx+(1โˆ’ฮป)y)โ‰คฮปf(x)+(1โˆ’ฮป)f(y)f(\lambda x + (1-\lambda)y) \leq \lambda f(x) + (1-\lambda) f(y)
Non-convex functions: Multi-modal, multiple local minima. Neural network training landscapes are non-convex. Generative model training is highly non-convex due to:
  • Neural network parameterizations
  • Min-max objectives (GANs)
  • Latent variable marginalization (VAEs)

1.5.2 Stochastic Gradient Descent (SGD)

ฮธt+1=ฮธtโˆ’ฮทtโˆ‡ฮธLB(ฮธt)\theta_{t+1} = \theta_t - \eta_t \nabla_\theta \mathcal{L}_B(\theta_t)
where LB\mathcal{L}_B is the loss computed on a mini-batch BB. Why SGD works for generative models:
  • Efficient for large datasets (compute gradient on subset)
  • Noise helps escape shallow local minima
  • Generalizes better than full-batch gradient descent

1.5.3 Adam Optimizer

Adam combines momentum + adaptive learning rates:
mt=ฮฒ1mtโˆ’1+(1โˆ’ฮฒ1)gt(momentum)m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t \quad \text{(momentum)} vt=ฮฒ2vtโˆ’1+(1โˆ’ฮฒ2)gt2(RMS)v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2 \quad \text{(RMS)} m^t=mt1โˆ’ฮฒ1t,v^t=vt1โˆ’ฮฒ2t(biasย correction)\hat{m}_t = \frac{m_t}{1-\beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1-\beta_2^t} \quad \text{(bias correction)} ฮธt+1=ฮธtโˆ’ฮทv^t+ฯตm^t\theta_{t+1} = \theta_t - \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon} \hat{m}_t
Recommended defaults for generative models:
  • ฮฒ1=0.9,ฮฒ2=0.999,ฮท=1eโˆ’4\beta_1 = 0.9, \beta_2 = 0.999, \eta = 1e-4 (GANs), 3eโˆ’43e-4 (diffusion models)
  • Lower learning rates for discriminator/critic (two-timescale update rule)

1.6 The ELBO: Full Derivation

The Evidence Lower Bound (ELBO) is the central objective for VAEs and many other generative models.

1.6.1 The Problem

We have a latent variable model:
pฮธ(x)=โˆซpฮธ(xโˆฃz)p(z)dzp_\theta(x) = \int p_\theta(x|z) p(z) dz
This integral is intractable for complex models (neural network decoders) because we'd need to integrate over all possible zz.

1.6.2 Introducing the Variational Distribution

We introduce an approximate posterior qฯ•(zโˆฃx)q_\phi(z|x) (the encoder). The goal: make qฯ•(zโˆฃx)q_\phi(z|x) close to the true posterior pฮธ(zโˆฃx)p_\theta(z|x).

1.6.3 Derivation Step 1: Start with Log-Likelihood

logโกpฮธ(x)=logโกโˆซpฮธ(xโˆฃz)p(z)dz\log p_\theta(x) = \log \int p_\theta(x|z) p(z) dz

1.6.4 Derivation Step 2: Multiply by 1 = q/q

logโกpฮธ(x)=logโกโˆซqฯ•(zโˆฃx)โ‹…pฮธ(xโˆฃz)p(z)qฯ•(zโˆฃx)dz\log p_\theta(x) = \log \int q_\phi(z|x) \cdot \frac{p_\theta(x|z) p(z)}{q_\phi(z|x)} dz

1.6.5 Derivation Step 3: Apply Jensen's Inequality

For a concave function (log is concave), logโก(E[y])โ‰ฅE[logโก(y)]\log(\mathbb{E}[y]) \geq \mathbb{E}[\log(y)]:
logโกpฮธ(x)=logโกEzโˆผqฯ•(zโˆฃx)[pฮธ(xโˆฃz)p(z)qฯ•(zโˆฃx)]โ‰ฅEzโˆผqฯ•(zโˆฃx)[logโกpฮธ(xโˆฃz)p(z)qฯ•(zโˆฃx)]\log p_\theta(x) = \log \mathbb{E}_{z \sim q_\phi(z|x)} \left[ \frac{p_\theta(x|z) p(z)}{q_\phi(z|x)} \right] \geq \mathbb{E}_{z \sim q_\phi(z|x)} \left[ \log \frac{p_\theta(x|z) p(z)}{q_\phi(z|x)} \right]
The RHS is the ELBO, denoted L(x;ฮธ,ฯ•)\mathcal{L}(x; \theta, \phi).

1.6.6 Derivation Step 4: Rearrange the ELBO

L(x;ฮธ,ฯ•)=Eqฯ•(zโˆฃx)[logโกpฮธ(xโˆฃz)]+Eqฯ•(zโˆฃx)[logโกp(z)qฯ•(zโˆฃx)]\mathcal{L}(x; \theta, \phi) = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] + \mathbb{E}_{q_\phi(z|x)}\left[\log \frac{p(z)}{q_\phi(z|x)}\right] =Eqฯ•(zโˆฃx)[logโกpฮธ(xโˆฃz)]โŸReconstructionโˆ’DKL(qฯ•(zโˆฃx)โˆฅp(z))โŸKLย divergence= \underbrace{\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]}_{\text{Reconstruction}} - \underbrace{D_{KL}(q_\phi(z|x) \| p(z))}_{\text{KL divergence}}

1.6.7 Derivation Step 5: The Gap

The gap between the true log-likelihood and the ELBO is the KL divergence between qฯ•(zโˆฃx)q_\phi(z|x) and the true posterior pฮธ(zโˆฃx)p_\theta(z|x):
logโกpฮธ(x)โˆ’L(x;ฮธ,ฯ•)=DKL(qฯ•(zโˆฃx)โˆฅpฮธ(zโˆฃx))\log p_\theta(x) - \mathcal{L}(x; \theta, \phi) = D_{KL}(q_\phi(z|x) \| p_\theta(z|x))
Proof:
DKL(qฯ•(zโˆฃx)โˆฅpฮธ(zโˆฃx))=Eqฯ•(zโˆฃx)[logโกqฯ•(zโˆฃx)pฮธ(zโˆฃx)]D_{KL}(q_\phi(z|x) \| p_\theta(z|x)) = \mathbb{E}_{q_\phi(z|x)}\left[\log \frac{q_\phi(z|x)}{p_\theta(z|x)}\right] =Eqฯ•(zโˆฃx)[logโกqฯ•(zโˆฃx)pฮธ(xโˆฃz)p(z)/pฮธ(x)]= \mathbb{E}_{q_\phi(z|x)}\left[\log \frac{q_\phi(z|x)}{p_\theta(x|z) p(z) / p_\theta(x)}\right] =Eqฯ•(zโˆฃx)[logโกqฯ•(zโˆฃx)โ‹…pฮธ(x)pฮธ(xโˆฃz)p(z)]= \mathbb{E}_{q_\phi(z|x)}\left[\log \frac{q_\phi(z|x) \cdot p_\theta(x)}{p_\theta(x|z) p(z)}\right] =logโกpฮธ(x)+Eqฯ•(zโˆฃx)[logโกqฯ•(zโˆฃx)pฮธ(xโˆฃz)p(z)]= \log p_\theta(x) + \mathbb{E}_{q_\phi(z|x)}\left[\log \frac{q_\phi(z|x)}{p_\theta(x|z) p(z)}\right] =logโกpฮธ(x)โˆ’Eqฯ•(zโˆฃx)[logโกpฮธ(xโˆฃz)]โˆ’DKL(qฯ•(zโˆฃx)โˆฅp(z))= \log p_\theta(x) - \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - D_{KL}(q_\phi(z|x) \| p(z)) =logโกpฮธ(x)โˆ’L(x;ฮธ,ฯ•)= \log p_\theta(x) - \mathcal{L}(x; \theta, \phi)

1.6.8 The Reparameterization Trick

The ELBO requires gradient through Ezโˆผqฯ•(zโˆฃx)[logโกpฮธ(xโˆฃz)]\mathbb{E}_{z \sim q_\phi(z|x)}[\log p_\theta(x|z)], but zz is sampled stochastically. Since sampling is non-differentiable, we reparameterize:
z=ฮผฯ•(x)+ฯƒฯ•(x)โ‹…ฯต,ฯตโˆผN(0,I)z = \mu_\phi(x) + \sigma_\phi(x) \cdot \epsilon, \quad \epsilon \sim \mathcal{N}(0, I)
Now the expectation is over ฯต\epsilon (which doesn't depend on ฯ•\phi), and the gradient can flow through ฮผฯ•\mu_\phi and ฯƒฯ•\sigma_\phi:
โˆ‡ฯ•Ezโˆผqฯ•[logโกpฮธ(xโˆฃz)]=EฯตโˆผN(0,I)[โˆ‡ฯ•logโกpฮธ(xโˆฃฮผฯ•(x)+ฯƒฯ•(x)โ‹…ฯต)]\nabla_\phi \mathbb{E}_{z \sim q_\phi}[\log p_\theta(x|z)] = \mathbb{E}_{\epsilon \sim \mathcal{N}(0, I)}[\nabla_\phi \log p_\theta(x | \mu_\phi(x) + \sigma_\phi(x) \cdot \epsilon)]

1.7 Why This Matters

The mathematical tools covered here are used throughout generative AI:
  • SVD: PCA for data preprocessing, analyzing latent space structure, low-rank approximations (LoRA) for fine-tuning
  • Eigendecomposition: Understanding covariance structure of features, FID computation
  • Matrix calculus: Every gradient-based training loop
  • Optimization: Choosing optimizers (Adam vs SGD), learning rate schedules, warmup
  • ELBO: Foundations of VAEs, diffusion models (which can be seen as hierarchical VAEs)

2. ๐Ÿ“ Key Formulas / Concepts

ConceptFormulaApplication in GenAI
EigendecompositionA=Vฮ›VTA = V \Lambda V^T (symmetric)PCA, covariance analysis
SVDA=UฮฃVTA = U \Sigma V^TLow-rank approximations, latent structure
Quadratic gradientโˆ‡x(xTAx)=2Ax\nabla_x (x^T A x) = 2Ax (A symmetric)Gaussian gradient computation
ELBO$\mathcal{L} = \mathbb{E}q[\log p\theta(xz)] - D_{KL}(q_\phi(z
Reparameterizationz=ฮผ+ฯƒโ‹…ฯตz = \mu + \sigma \cdot \epsilonDifferentiable sampling
Adam updateฮธt+1=ฮธtโˆ’ฮทmtvt+ฯต\theta_{t+1} = \theta_t - \eta \frac{m_t}{\sqrt{v_t}+\epsilon}Default optimizer for gen models

3. โš ๏ธ Common Pitfalls

Pitfall 1: Confusing SVD and Eigendecomposition

Mistake: Using eigendecomposition when SVD is needed (for non-square matrices). Why: Eigendecomposition only works for square matrices. SVD works for any matrix. Correct approach: Use eigendecomposition for analyzing square symmetric matrices (covariance matrices, Gram matrices). Use SVD for data matrices, weight matrices, and any non-square case.

Pitfall 2: Forgetting the Reparameterization Trick Isn't Always Applicable

Mistake: Trying to apply the reparameterization trick to discrete latent variables. Why: The reparameterization trick requires the sampling distribution to be reparameterizable as a deterministic function of a noise source. For discrete distributions (e.g., categorical), zz takes discrete values and can't be written as ฮผ+ฯƒโ‹…ฯต\mu + \sigma \cdot \epsilon with ฯตโˆผN(0,I)\epsilon \sim \mathcal{N}(0, I). Correct approach: Use the Gumbel-Softmax trick (continuous relaxation of discrete sampling) or REINFORCE gradient estimator for discrete latents.

Pitfall 3: Ignoring the KL Vanishing Problem

Mistake: Training a VAE and wondering why the latent variables aren't used. Why: The KL term in the ELBO encourages qฯ•(zโˆฃx)q_\phi(z|x) to match the prior p(z)p(z). If the decoder is powerful enough, the model can "ignore" zz and set qฯ•(zโˆฃx)=p(z)q_\phi(z|x) = p(z), resulting in DKL=0D_{KL}=0 but no meaningful latent representation. Correct approach: Monitor KL divergence during training. If it drops to near zero, use KL annealing, free bits (a minimum KL per dimension), or reduce decoder capacity.

Pitfall 4: Using SGD when Adam is Needed

Mistake: Using vanilla SGD for GAN or VAE training with poor results. Why: Generative model losses are highly non-convex and have varying gradient scales across parameters. SGD's fixed learning rate per parameter struggles. Adam's adaptive learning rates handle the varying scales better. Correct approach: Use Adam for most generative models. For GANs, use two-timescale Adam (different learning rates for generator and discriminator).

Pitfall 5: Mishandling Matrix Calculus Gradients

Mistake: Computing matrix gradients as scalars without respecting transpose relationships. Why: โˆ‚fโˆ‚W\frac{\partial f}{\partial W} has the same dimensions as WW, but the chain rule for matrix functions requires attention to transposition. Correct approach: Use dimensional analysis. If f=โˆฅWxโˆ’yโˆฅ2f = \|Wx - y\|^2:
  • ff is scalar
  • WW is doutร—dind_{out} \times d_{in}
  • โˆ‚fโˆ‚W=2(Wxโˆ’y)xT\frac{\partial f}{\partial W} = 2(Wx - y)x^T (check: doutร—dind_{out} \times d_{in} โœ“)

4. ๐Ÿ“ Practice Questions

**Q1: Compute the eigenvalues of
>A=[2112]>> A = \begin{bmatrix} 2 & 1 \\ 1 & 2 \end{bmatrix} >
and explain what they mean geometrically.**
Step 1: detโก(Aโˆ’ฮปI)=(2โˆ’ฮป)2โˆ’1=ฮป2โˆ’4ฮป+3=0\det(A - \lambda I) = (2-\lambda)^2 - 1 = \lambda^2 - 4\lambda + 3 = 0
Step 2: ฮป=4ยฑ16โˆ’122=4ยฑ22\lambda = \frac{4 \pm \sqrt{16 - 12}}{2} = \frac{4 \pm 2}{2}
Step 3: ฮป1=3\lambda_1 = 3, ฮป2=1\lambda_2 = 1
Geometric interpretation: AA stretches space by a factor of 3 along eigenvectors corresponding to ฮป1\lambda_1 and by 1 (no stretch) along eigenvectors corresponding to ฮป2\lambda_2.
Eigenvectors:
  • ฮป1=3\lambda_1 = 3: v1=[1,1]T/2v_1 = [1, 1]^T/\sqrt{2} (direction of equal increase)
  • ฮป2=1\lambda_2 = 1: v2=[1,โˆ’1]T/2v_2 = [1, -1]^T/\sqrt{2} (direction of difference)
>A=[2112]=Vฮ›VT=12[111โˆ’1][3001][111โˆ’1]>> A = \begin{bmatrix} 2 & 1 \\ 1 & 2 \end{bmatrix} = V \Lambda V^T = \frac{1}{2}\begin{bmatrix} 1 & 1 \\ 1 & -1 \end{bmatrix} \begin{bmatrix} 3 & 0 \\ 0 & 1 \end{bmatrix} \begin{bmatrix} 1 & 1 \\ 1 & -1 \end{bmatrix} >
**Q2: Compute the SVD of
>A=[100020]>> A = \begin{bmatrix} 1 & 0 & 0 \\ 0 & 2 & 0 \end{bmatrix} >
.**
Step 1: AA is 2ร—32 \times 3, so UU is 2ร—22 \times 2, ฮฃ\Sigma is 2ร—32 \times 3, VV is 3ร—33 \times 3.
Step 2:
>AAT=[1004]>> AA^T = \begin{bmatrix} 1 & 0 \\ 0 & 4 \end{bmatrix} >
, eigenvalues: ฮป1=4\lambda_1 = 4, ฮป2=1\lambda_2 = 1.
Singular values: ฯƒ1=2\sigma_1 = 2, ฯƒ2=1\sigma_2 = 1.
Step 3: UU from eigenvectors of AATAA^T: u1=[0,1]Tu_1 = [0, 1]^T, u2=[1,0]Tu_2 = [1, 0]^T.
Step 4: VV from V=ATUฮฃโˆ’1V = A^T U \Sigma^{-1} (more carefully):
ฮฃ=[diag(2,1),02ร—1]\Sigma = [\text{diag}(2, 1), 0_{2 \times 1}]
>v1=ATu1ฯƒ1=12[100200][01]=12[020]=[010]>> v_1 = \frac{A^T u_1}{\sigma_1} = \frac{1}{2} \begin{bmatrix} 1 & 0 \\ 0 & 2 \\ 0 & 0 \end{bmatrix} \begin{bmatrix} 0 \\ 1 \end{bmatrix} = \frac{1}{2} \begin{bmatrix} 0 \\ 2 \\ 0 \end{bmatrix} = \begin{bmatrix} 0 \\ 1 \\ 0 \end{bmatrix} >
>v2=ATu2ฯƒ2=11[100200][10]=[100]>> v_2 = \frac{A^T u_2}{\sigma_2} = \frac{1}{1} \begin{bmatrix} 1 & 0 \\ 0 & 2 \\ 0 & 0 \end{bmatrix} \begin{bmatrix} 1 \\ 0 \end{bmatrix} = \begin{bmatrix} 1 \\ 0 \\ 0 \end{bmatrix} >
v3v_3 is in the nullspace of AA (since Av3=0A v_3 = 0): v3=[0,0,1]Tv_3 = [0, 0, 1]^T.
Step 5: Verify:
>UฮฃVT=[0110][200010][010100001]=[100020]=A>> U \Sigma V^T = \begin{bmatrix} 0 & 1 \\ 1 & 0 \end{bmatrix} \begin{bmatrix} 2 & 0 & 0 \\ 0 & 1 & 0 \end{bmatrix} \begin{bmatrix} 0 & 1 & 0 \\ 1 & 0 & 0 \\ 0 & 0 & 1 \end{bmatrix} = \begin{bmatrix} 1 & 0 & 0 \\ 0 & 2 & 0 \end{bmatrix} = A >
Q3: Compute โˆ‡xf\nabla_x f for f(x)=โˆฅAxโˆ’bโˆฅ22f(x) = \|Ax - b\|_2^2.
Step 1: Expand: f=(Axโˆ’b)T(Axโˆ’b)=xTATAxโˆ’2bTAx+bTbf = (Ax - b)^T (Ax - b) = x^T A^T A x - 2b^T A x + b^T b
Step 2: Differentiate each term:
  • โˆ‚โˆ‚x(xTATAx)=2ATAx\frac{\partial}{\partial x} (x^T A^T A x) = 2 A^T A x
  • โˆ‚โˆ‚x(โˆ’2bTAx)=โˆ’2ATb\frac{\partial}{\partial x} (-2b^T A x) = -2 A^T b
  • โˆ‚โˆ‚x(bTb)=0\frac{\partial}{\partial x} (b^T b) = 0
Step 3: โˆ‡xf=2ATAxโˆ’2ATb=2AT(Axโˆ’b)\nabla_x f = 2 A^T A x - 2 A^T b = 2 A^T (Ax - b) Q4: Starting from logโกp(x)\log p(x), derive the ELBO and show that optimizing it maximizes a lower bound on the log-likelihood.
The derivation is in Section 1.6 above. The key steps:
  1. logโกp(x)=logโกโˆซp(xโˆฃz)p(z)dz\log p(x) = \log \int p(x|z) p(z) dz
  2. Introduce q(zโˆฃx)q(z|x): logโกp(x)=logโกโˆซq(zโˆฃx)p(xโˆฃz)p(z)q(zโˆฃx)dz\log p(x) = \log \int q(z|x) \frac{p(x|z)p(z)}{q(z|x)} dz
  3. Jensen's inequality: logโกEq[f]โ‰ฅEq[logโกf]\log \mathbb{E}_q[f] \geq \mathbb{E}_q[\log f]
  4. Result: logโกp(x)โ‰ฅEq[logโกp(xโˆฃz)]โˆ’DKL(q(zโˆฃx)โˆฅp(z))=ELBO\log p(x) \geq \mathbb{E}_q[\log p(x|z)] - D_{KL}(q(z|x) \| p(z)) = \text{ELBO}
Optimizing ฮธ,ฯ•\theta, \phi to maximize the ELBO:
  • Increases the log-likelihood (by at least as much as the ELBO increases)
  • Decreases DKL(qฯ•(zโˆฃx)โˆฅpฮธ(zโˆฃx))D_{KL}(q_\phi(z|x) \| p_\theta(z|x)) (the approximation gap)
The tightness depends on how well qฯ•(zโˆฃx)q_\phi(z|x) approximates the true posterior pฮธ(zโˆฃx)p_\theta(z|x). Q5: In the reparameterization trick, why can't we just differentiate through the sampling operation directly?
The sampling operation zโˆผqฯ•(zโˆฃx)z \sim q_\phi(z|x) is a stochastic process. The gradient of this operation w.r.t. ฯ•\phi would require differentiating through a random number generator, which is non-differentiable.
Concretely, z=sample(ฮผฯ•,ฯƒฯ•)z = \text{sample}(\mu_\phi, \sigma_\phi) involves calling a random number generator to produce ฯตโˆผN(0,1)\epsilon \sim \mathcal{N}(0, 1), then computing z=ฮผฯ•+ฯƒฯ•ฯตz = \mu_\phi + \sigma_\phi \epsilon. The RNG call has no gradient (the operation is discrete โ€” a number is drawn; there's no continuous path for gradient flow).
The reparameterization trick moves the stochasticity to an independent noise source ฯต\epsilon:
z=ฮผฯ•(x)+ฯƒฯ•(x)โ‹…ฯต,ฯตโˆผN(0,1)z = \mu_\phi(x) + \sigma_\phi(x) \cdot \epsilon, \quad \epsilon \sim \mathcal{N}(0, 1)
Now:
  • The gradient can flow through ฮผฯ•\mu_\phi and ฯƒฯ•\sigma_\phi (differentiable operations)
  • ฯต\epsilon is independent of ฯ•\phi, so no gradient through sampling
  • The expectation is still valid because N(zโˆฃฮผ,ฯƒ2)dz=N(ฯตโˆฃ0,1)dฯต\mathcal{N}(z|\mu, \sigma^2) dz = \mathcal{N}(\epsilon|0, 1) d\epsilon Q6: For a covariance matrix ฮฃ\Sigma, explain why ฮฃ\Sigma must be positive semidefinite and how eigendecomposition reveals this.
A covariance matrix must be positive semidefinite (PSD) because for any vector vv:
vTฮฃv=vTE[(Xโˆ’ฮผ)(Xโˆ’ฮผ)T]v=E[vT(Xโˆ’ฮผ)(Xโˆ’ฮผ)Tv]=E[(vT(Xโˆ’ฮผ))2]โ‰ฅ0v^T \Sigma v = v^T \mathbb{E}[(X - \mu)(X - \mu)^T] v = \mathbb{E}[v^T (X - \mu)(X - \mu)^T v] = \mathbb{E}[(v^T (X - \mu))^2] \geq 0
The variance of any linear combination is non-negative.
By eigendecomposition, ฮฃ=Vฮ›VT\Sigma = V \Lambda V^T. PSD means all eigenvalues ฮปiโ‰ฅ0\lambda_i \geq 0. This is evident because:
  • If any ฮปi<0\lambda_i < 0, then viTฮฃvi=ฮปi<0v_i^T \Sigma v_i = \lambda_i < 0 (contradiction)
  • The singular values from SVD of the data matrix are ฮปi\sqrt{\lambda_i}, confirming non-negativity
This PSD property ensures:
  • FID computation (matrix square root) is valid
  • KL divergence between Gaussians is well-defined
  • Cholesky decomposition for efficient sampling exists Q7: A VAE uses Adam with default parameters (ฮท=0.001\eta=0.001, ฮฒ1=0.9\beta_1=0.9, ฮฒ2=0.999\beta_2=0.999). After 100 epochs, the KL term is 0.01 (near zero). Diagnose the problem and propose a fix.
Diagnosis: KL vanishing. The model is ignoring the latent variable zz and functioning as an autoencoder with no regularization. The decoder is powerful enough to reconstruct xx without zz, so the KL term forces q(zโˆฃx)โ†’p(z)=N(0,I)q(z|x) \to p(z) = \mathcal{N}(0, I), giving DKLโ‰ˆ0D_{KL} \approx 0.
Fixes (in order of recommendation):
  1. KL annealing: Start with ฮฒ=0\beta = 0 in L=Reconโˆ’ฮฒโ‹…KL\mathcal{L} = \text{Recon} - \beta \cdot KL, gradually increase ฮฒ\beta from 0 to 1 over training. This lets the model learn meaningful latent representations before the KL penalty kicks in.
  2. Free bits: Set a minimum KL per dimension: KLmin=maxโก(ฮฒโ‹…KL,ฮป)KL_{\text{min}} = \max(\beta \cdot KL, \lambda). This ensures each dimension carries at least ฮป\lambda nats of information.
  3. Reduce decoder capacity: A weaker decoder (fewer layers, smaller hidden size) forces the model to use the latent code.
  4. Increase latent dimension: More dimensions give the model more capacity to encode information before the KL penalty forces them to the prior.
Q8: Show that โˆ‡ฮธExโˆผpฮธ[f(x)]=Exโˆผpฮธ[f(x)โˆ‡ฮธlogโกpฮธ(x)]\nabla_\theta \mathbb{E}_{x \sim p_\theta}[f(x)] = \mathbb{E}_{x \sim p_\theta}[f(x) \nabla_\theta \log p_\theta(x)] (the REINFORCE/score function gradient).
Step 1: Write the expectation as an integral:
โˆ‡ฮธExโˆผpฮธ[f(x)]=โˆ‡ฮธโˆซf(x)pฮธ(x)dx\nabla_\theta \mathbb{E}_{x \sim p_\theta}[f(x)] = \nabla_\theta \int f(x) p_\theta(x) dx
Step 2: Swap gradient and integral (under smoothness conditions):
=โˆซf(x)โˆ‡ฮธpฮธ(x)dx= \int f(x) \nabla_\theta p_\theta(x) dx
Step 3: Use the log-derivative trick: โˆ‡ฮธpฮธ(x)=pฮธ(x)โˆ‡ฮธlogโกpฮธ(x)\nabla_\theta p_\theta(x) = p_\theta(x) \nabla_\theta \log p_\theta(x)
=โˆซf(x)pฮธ(x)โˆ‡ฮธlogโกpฮธ(x)dx= \int f(x) p_\theta(x) \nabla_\theta \log p_\theta(x) dx
Step 4: Recognize as an expectation:
=Exโˆผpฮธ[f(x)โˆ‡ฮธlogโกpฮธ(x)]= \mathbb{E}_{x \sim p_\theta}[f(x) \nabla_\theta \log p_\theta(x)]
This is the REINFORCE gradient estimator or score function estimator. It allows gradient estimation through non-differentiable sampling by using logโกpฮธ(x)\log p_\theta(x) instead of โˆ‡ฮธx\nabla_\theta x.
Comparison with reparameterization:
  • Reparameterization: Lower variance, but requires continuous xx and differentiable ff
  • REINFORCE: Higher variance, but works for discrete xx and non-differentiable ff Q9: For a multivariate Gaussian N(ฮผ,ฮฃ)\mathcal{N}(\mu, \Sigma), derive the gradient of the log-likelihood with respect to ฮผ\mu for a single sample xx.
Step 1: Write the log-likelihood:
logโกp(xโˆฃฮผ,ฮฃ)=โˆ’12(xโˆ’ฮผ)Tฮฃโˆ’1(xโˆ’ฮผ)โˆ’12logโกโˆฃฮฃโˆฃโˆ’D2logโก(2ฯ€)\log p(x|\mu, \Sigma) = -\frac{1}{2}(x-\mu)^T \Sigma^{-1} (x-\mu) - \frac{1}{2}\log|\Sigma| - \frac{D}{2}\log(2\pi)
Step 2: Only the first term depends on ฮผ\mu:
โˆ‡ฮผlogโกp(xโˆฃฮผ,ฮฃ)=โˆ’12โˆ‡ฮผ[(xโˆ’ฮผ)Tฮฃโˆ’1(xโˆ’ฮผ)]\nabla_\mu \log p(x|\mu, \Sigma) = -\frac{1}{2} \nabla_\mu [(x-\mu)^T \Sigma^{-1} (x-\mu)]
Step 3: Let y=xโˆ’ฮผy = x - \mu. Then โˆ‚yโˆ‚ฮผ=โˆ’I\frac{\partial y}{\partial \mu} = -I.
Step 4: โˆ‚โˆ‚y(yTฮฃโˆ’1y)=2ฮฃโˆ’1y\frac{\partial}{\partial y} (y^T \Sigma^{-1} y) = 2\Sigma^{-1} y (since ฮฃโˆ’1\Sigma^{-1} is symmetric)
Step 5: Chain rule:
โˆ‡ฮผlogโกp=โˆ’12โ‹…โˆ‚(yTฮฃโˆ’1y)โˆ‚yโ‹…โˆ‚yโˆ‚ฮผ=โˆ’12โ‹…(2ฮฃโˆ’1y)Tโ‹…(โˆ’I)=yTฮฃโˆ’1=(xโˆ’ฮผ)Tฮฃโˆ’1\nabla_\mu \log p = -\frac{1}{2} \cdot \frac{\partial (y^T \Sigma^{-1} y)}{\partial y} \cdot \frac{\partial y}{\partial \mu} = -\frac{1}{2} \cdot (2\Sigma^{-1} y)^T \cdot (-I) = y^T \Sigma^{-1} = (x-\mu)^T \Sigma^{-1}
In column-vector convention: โˆ‡ฮผlogโกp=ฮฃโˆ’1(xโˆ’ฮผ)\nabla_\mu \log p = \Sigma^{-1} (x - \mu).
The MLE for ฮผ\mu sets this gradient to zero: ฮฃโˆ’1(ฮผ^โˆ’xavg)=0โ€…โ€ŠโŸนโ€…โ€Šฮผ^=1Nโˆ‘xi\Sigma^{-1} (\hat{\mu} - x_{\text{avg}}) = 0 \implies \hat{\mu} = \frac{1}{N}\sum x_i, confirming that the MLE for the mean is the sample mean. Q10: Explain why Adam's bias correction is important in the first few iterations of VAE training.
Without bias correction, Adam's moment estimates are biased toward zero because they're initialized as m0=0,v0=0m_0 = 0, v_0 = 0:
mt=ฮฒ1mtโˆ’1+(1โˆ’ฮฒ1)gtm_t = \beta_1 m_{t-1} + (1-\beta_1) g_t
In early iterations, mtm_t is "warmed up" from 0. For ฮฒ1=0.9\beta_1 = 0.9:
  • t=1t=1: m1=0.9โ‹…0+0.1โ‹…g1=0.1g1m_1 = 0.9 \cdot 0 + 0.1 \cdot g_1 = 0.1 g_1 (10% of true gradient)
  • t=5t=5: m5โ‰ˆ0.41gavgm_5 \approx 0.41 g_{\text{avg}} (still biased)
  • t=20t=20: m20โ‰ˆ0.88gavgm_{20} \approx 0.88 g_{\text{avg}} (approaching true)
Bias correction divides by 1โˆ’ฮฒ1t1 - \beta_1^t:
m^t=mt1โˆ’ฮฒ1t\hat{m}_t = \frac{m_t}{1 - \beta_1^t}
  • t=1t=1: m^1=0.1g1/(1โˆ’0.9)=g1\hat{m}_1 = 0.1 g_1 / (1 - 0.9) = g_1 โœ“
  • t=5t=5: m^5โ‰ˆ0.41gavg/(1โˆ’0.95)=0.41/0.41โ‰ˆgavg\hat{m}_5 \approx 0.41 g_{\text{avg}} / (1 - 0.9^5) = 0.41 / 0.41 \approx g_{\text{avg}} โœ“
Without bias correction, the first steps would be too small, slowing convergence โ€” particularly noticeable in VAE training where initial gradients guide the model toward meaningful latent representations. Q11: Derive the gradient of DKL(qฯ•(zโˆฃx)โˆฅp(z))D_{KL}(q_\phi(z|x) \| p(z)) with respect to ฯ•\phi where qฯ•(zโˆฃx)=N(ฮผฯ•,ฯƒฯ•2)q_\phi(z|x) = \mathcal{N}(\mu_\phi, \sigma_\phi^2) and p(z)=N(0,1)p(z) = \mathcal{N}(0, 1).
Step 1: Write the KL in closed form for Gaussians:
DKL(N(ฮผ,ฯƒ2)โˆฅN(0,1))=logโก1ฯƒ+ฯƒ2+ฮผ22โˆ’12D_{KL}(\mathcal{N}(\mu, \sigma^2) \| \mathcal{N}(0, 1)) = \log\frac{1}{\sigma} + \frac{\sigma^2 + \mu^2}{2} - \frac{1}{2}
=โˆ’logโกฯƒ+ฯƒ2+ฮผ22โˆ’12= -\log\sigma + \frac{\sigma^2 + \mu^2}{2} - \frac{1}{2}
Step 2: Differentiate w.r.t. ฮผ\mu:
โˆ‚KLโˆ‚ฮผ=ฮผ\frac{\partial KL}{\partial \mu} = \mu
Step 3: Differentiate w.r.t. ฯƒ\sigma (using ฯƒ\sigma as the standard deviation):
โˆ‚KLโˆ‚ฯƒ=โˆ’1ฯƒ+ฯƒ=ฯƒโˆ’1ฯƒ=ฯƒ2โˆ’1ฯƒ\frac{\partial KL}{\partial \sigma} = -\frac{1}{\sigma} + \sigma = \sigma - \frac{1}{\sigma} = \frac{\sigma^2 - 1}{\sigma}
Step 4: In practice, we parameterize logโกฯƒ2\log\sigma^2 (log-variance) for numerical stability. Let s=logโกฯƒ2s = \log\sigma^2, so ฯƒ=es/2\sigma = e^{s/2}:
โˆ‚KLโˆ‚s=โˆ‚KLโˆ‚ฯƒโ‹…โˆ‚ฯƒโˆ‚s=(ฯƒโˆ’1ฯƒ)โ‹…es/22=(es/2โˆ’eโˆ’s/2)โ‹…es/22\frac{\partial KL}{\partial s} = \frac{\partial KL}{\partial \sigma} \cdot \frac{\partial \sigma}{\partial s} = \left(\sigma - \frac{1}{\sigma}\right) \cdot \frac{e^{s/2}}{2} = \left(e^{s/2} - e^{-s/2}\right) \cdot \frac{e^{s/2}}{2}
=esโˆ’12=ฯƒ2โˆ’12= \frac{e^s - 1}{2} = \frac{\sigma^2 - 1}{2}
The gradient tells us:
  • If ฮผ>0\mu > 0, push ฮผ\mu toward zero (reduce mean)
  • If ฯƒ>1\sigma > 1, push ฯƒ\sigma toward 1 (reduce variance)
  • If ฯƒ<1\sigma < 1, push ฯƒ\sigma toward 1 (increase variance)
The KL pushes the posterior toward a unit Gaussian โ€” neither too wide nor too narrow. Q12: In diffusion models, the training loss L=Et,x0,ฯต[โˆฅฯตโˆ’ฯตฮธ(xt,t)โˆฅ2]\mathcal{L} = \mathbb{E}_{t, x_0, \epsilon}[\|\epsilon - \epsilon_\theta(x_t, t)\|^2] resembles a denoising objective. Show how this relates to the ELBO.
The diffusion model's variational bound can be written as:
Lvb=โˆ‘t=1TEx0,ฯตt[DKL(q(xtโˆ’1โˆฃxt,x0)โˆฅpฮธ(xtโˆ’1โˆฃxt))]\mathcal{L}_{vb} = \sum_{t=1}^T \mathbb{E}_{x_0, \epsilon_t}[D_{KL}(q(x_{t-1}|x_t, x_0) \| p_\theta(x_{t-1}|x_t))]
Each KL term is between two Gaussians:
  • q(xtโˆ’1โˆฃxt,x0)=N(ฮผ~t(xt,x0),ฮฒ~tI)q(x_{t-1}|x_t, x_0) = \mathcal{N}(\tilde{\mu}_t(x_t, x_0), \tilde{\beta}_t I) (the forward posterior, known)
  • pฮธ(xtโˆ’1โˆฃxt)=N(ฮผฮธ(xt,t),ฯƒt2I)p_\theta(x_{t-1}|x_t) = \mathcal{N}(\mu_\theta(x_t, t), \sigma_t^2 I) (the model, learned)
Using the closed-form KL for Gaussians:
DKL(qโˆฅpฮธ)=12ฯƒt2โˆฅฮผ~t(xt,x0)โˆ’ฮผฮธ(xt,t)โˆฅ2+constD_{KL}(q \| p_\theta) = \frac{1}{2\sigma_t^2} \|\tilde{\mu}_t(x_t, x_0) - \mu_\theta(x_t, t)\|^2 + \text{const}
We parameterize ฮผฮธ\mu_\theta as:
ฮผฮธ(xt,t)=1ฮฑt(xtโˆ’ฮฒt1โˆ’ฮฑห‰tฯตฮธ(xt,t))\mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}} \left(x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon_\theta(x_t, t)\right)
And the forward posterior mean is:
ฮผ~t(xt,x0)=1ฮฑt(xtโˆ’ฮฒt1โˆ’ฮฑห‰tฯตt)\tilde{\mu}_t(x_t, x_0) = \frac{1}{\sqrt{\alpha_t}} \left(x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon_t\right)
Substituting, the KL simplifies to:
DKL(qโˆฅpฮธ)โˆโˆฅฯตtโˆ’ฯตฮธ(xt,t)โˆฅ2D_{KL}(q \| p_\theta) \propto \|\epsilon_t - \epsilon_\theta(x_t, t)\|^2
Thus, the diffusion training loss โˆฅฯตโˆ’ฯตฮธ(xt,t)โˆฅ2\|\epsilon - \epsilon_\theta(x_t, t)\|^2 is equivalent to minimizing the KL divergence between the true denoising step and the model's prediction at each timestep โ€” which is exactly the ELBO for diffusion models. This demonstrates the deep connection between denoising score matching and variational inference.

5. ๐Ÿ”— Cross-References

  • All BSDA5002 topics: This math is used throughout
  • Information Theory (Week 9) (../week09/09-information-theory.md) โ€” KL divergence, entropy
  • VAEs (Week 3) (../week03/03-vaes.md) โ€” ELBO in practice
  • External: Matrix Cookbook by Petersen & Pedersen โ€” Comprehensive matrix calculus reference
  • External: "Deep Learning" by Goodfellow, Bengio, Courville โ€” Chapters on optimization and linear algebra Join Discord PreviousInformation TheoryNextBSDA5002 โ€” Generative AI Foundations
Document outline

Keep your place and jump directly to a heading.

Table of contents

No note selected

Open a chapter to see its outline, backlinks, and sources here.