Multivariate Gaussian distributions

A generative model is a distribution we can sample from and, ideally, evaluate (What generative modeling is). Before neural networks enter, the course works with the simplest distribution that is still interesting in several dimensions: the multivariate Gaussian. It is the running example of the first notes, where it is both the target we draw data from and the model we fit to that data. It also comes back in almost every later model: the prior of a VAE, the noise of a diffusion model, the base distribution of a normalizing flow and the feature statistics behind FID.

The opening block of the course has three steps:

1. generate data from a known distribution (this note);
2. measure how far two distributions are from each other (Entropy, cross-entropy and KL divergence);
3. fit a parametrized 2-D distribution to the data by gradient descent on such a distance (Fitting distributions with KL).

Notation. Capital letters such as $X$ and $Z$ denote random vectors; bold lowercase letters such as $\mathbf x$ and $\mathbf z$ denote points of the plane, the values those vectors take. Bold capitals such as $\boldsymbol\Sigma$, $\mathbf A$ and $\mathbf L$ are matrices, $\mathbf I$ is the identity matrix, $\mathbf A^{\top}$ the transpose and $|\boldsymbol\Sigma|$ the determinant. We write $X \sim \mathcal N(\mu, \sigma^2)$ for a univariate normal with mean $\mu$ and variance $\sigma^2$, and $X \sim \mathcal N(\boldsymbol\mu, \boldsymbol\Sigma)$ for a multivariate normal with mean vector $\boldsymbol\mu$ and covariance matrix $\boldsymbol\Sigma$: the second argument is always the variance or covariance, never the standard deviation. Densities are written $p$ and $q$; from the PyTorch section on, $P$ is the fixed target distribution with density $p$ and $Q$ the trainable model with density $q$.

The covariance matrix#

Take a random column vector in two dimensions and its mean:

$$X = \begin{pmatrix} X_1 \\ X_2 \end{pmatrix}, \qquad \boldsymbol{\mu} = \mathbb E[X] = \begin{pmatrix} \mu_{1} \\ \mu_{2} \end{pmatrix}.$$

Its spread is described by the $2 \times 2$ variance–covariance matrix $\boldsymbol\Sigma = \mathbb E\bigl[(X-\boldsymbol\mu)(X-\boldsymbol\mu)^{\top}\bigr]$:

$$\begin{aligned} \boldsymbol{\Sigma} &= \begin{pmatrix} \operatorname{Var}(X_1) & \operatorname{Cov}(X_1, X_2) \\ \operatorname{Cov}(X_1, X_2) & \operatorname{Var}(X_2) \end{pmatrix} \\[4pt] &= \begin{pmatrix} \sigma_{1}^2 & \rho\,\sigma_{1}\sigma_{2} \\ \rho\,\sigma_{1}\sigma_{2} & \sigma_{2}^2 \end{pmatrix}, \end{aligned}$$

where $\rho$ is the correlation coefficient between $X_1$ and $X_2$. The matrix is symmetric, because $\operatorname{Cov}(X_1, X_2) = \operatorname{Cov}(X_2, X_1)$. It is also positive semidefinite: for any fixed vector $\mathbf a$ the variance $\operatorname{Var}(\mathbf a^{\top}X) = \mathbf a^{\top}\boldsymbol\Sigma\,\mathbf a$ cannot be negative. Three free numbers hold every variance and covariance of the vector.

The example used throughout this note is a Gaussian centred at the origin with

$$\boldsymbol{\mu} = \begin{pmatrix} 0 \\ 0 \end{pmatrix}, \qquad \boldsymbol{\Sigma} = \begin{pmatrix} 1 & 0.5 \\ 0.5 & 1 \end{pmatrix}:$$

unit variances and correlation $\rho = 0.5$, so the cloud leans along the diagonal $x_1 = x_2$.

Why a full matrix and not just two variances?#

Two variances describe how far the cloud spreads along each coordinate axis. That is enough only when the coordinates are uncorrelated, and then the cloud is stretched along the axes. The off-diagonal entry records what the variances cannot: how the coordinates move together. It tilts the cloud.

In Figure 1 the right cloud is the left one rotated by 45°. The two clouds have the same eigenvalues (1.5 and 0.5), but both axis variances of the right cloud equal 1. Those are exactly the axis variances of the round $\mathcal N(\mathbf 0, \mathbf I)$ cloud, so two variances alone cannot tell the tilted cloud from a circle. Only $\operatorname{Cov}(X_1, X_2) = 0.5$ records the tilt.

Figure

A diagonal matrix only stretches along the axes; covariance tilts the cloud

The bivariate density#

When $\boldsymbol\Sigma$ is invertible (positive definite, not just semidefinite), the joint density on the whole plane $\mathbb R^2$ is

$$\begin{aligned} p(\mathbf x) &= \frac{1}{2\pi\sqrt{|\boldsymbol\Sigma|}} \\[4pt] &\times \exp\!\left(-\tfrac12\, (\mathbf{x}-\boldsymbol\mu)^{\top}\,\boldsymbol\Sigma^{-1}\,(\mathbf{x}-\boldsymbol\mu)\right), \end{aligned}$$

where $\mathbf x = (x_1, x_2)^{\top}$, $\boldsymbol\Sigma^{-1}$ is the inverse and $|\boldsymbol\Sigma|$ the determinant of the covariance matrix. The shapes in the quadratic form are

$$\underbrace{(\mathbf{x}-\boldsymbol\mu)^{\top}}_{1\times 2}\; \underbrace{\boldsymbol\Sigma^{-1}}_{2\times 2}\; \underbrace{(\mathbf{x}-\boldsymbol\mu)}_{2\times 1},$$

so the exponent is a scalar. In $d$ dimensions the normalizing constant becomes $(2\pi)^{-d/2}|\boldsymbol\Sigma|^{-1/2}$, and $d = 2$ gives the expression above (Bishop, 2006, §2.3).

The level sets of the density, where the squared Mahalanobis distance $(\mathbf{x}-\boldsymbol\mu)^{\top}\boldsymbol\Sigma^{-1}(\mathbf{x}-\boldsymbol\mu)$ equals $k^2$, are ellipses (Murphy, 2022, §3.2). Figures 1 and 3 draw them as the $1\sigma$ ($k = 1$) and $2\sigma$ ($k = 2$) contours. The probability inside such an ellipse follows from the map of the sampling section below. Writing $X = \boldsymbol\mu + \boldsymbol\Sigma^{1/2} Z$ with $Z \sim \mathcal N(\mathbf 0, \mathbf I)$, and using that $\boldsymbol\Sigma^{1/2}$ is symmetric with $\boldsymbol\Sigma^{1/2}\boldsymbol\Sigma^{-1}\boldsymbol\Sigma^{1/2} = \mathbf I$, the squared distance becomes $Z^{\top}Z = Z_1^2 + Z_2^2$, the squared length of standard noise. The standard density $\frac{1}{2\pi}e^{-r^2/2}$ depends only on the radius $r$, so in polar coordinates

$$\begin{aligned} \Pr\bigl(Z_1^2 + Z_2^2 \le k^2\bigr) &= \int_0^{k} r\,e^{-r^2/2}\,dr \\[4pt] &= 1 - e^{-k^2/2}. \end{aligned}$$

This is the CDF of the $\chi^2$ distribution with two degrees of freedom, evaluated at $k^2$. So the $1\sigma$ ellipse holds $1 - e^{-1/2} \approx 39\%$ of the probability and the $2\sigma$ ellipse $1 - e^{-2} \approx 86\%$, not the 68% and 95% of one-dimensional intervals.

For the running example

$$\begin{aligned} |\boldsymbol\Sigma| &= 1 - 0.5^2 = 0.75, \\[4pt] \boldsymbol\Sigma^{-1} &= \frac{1}{0.75}\begin{pmatrix} 1 & -0.5 \\ -0.5 & 1 \end{pmatrix} \\[4pt] &= \begin{pmatrix} 4/3 & -2/3 \\ -2/3 & 4/3 \end{pmatrix}, \end{aligned}$$

and the peak density is $p(\boldsymbol\mu) = 1/(2\pi\sqrt{0.75}) \approx 0.184$.

Sampling: from standard noise to any Gaussian#

In short. To sample $\mathcal N(\boldsymbol\mu, \boldsymbol\Sigma)$, draw $Z \sim \mathcal N(\mathbf 0, \mathbf I)$ and return $\boldsymbol\mu + \mathbf A Z$ for any matrix $\mathbf A$ with $\mathbf A\mathbf A^{\top} = \boldsymbol\Sigma$; in code $\mathbf A$ is the Cholesky factor. The rest of this section builds the square root and proves the recipe twice, by the change of variables and by the moments. Readers who know affine maps of Gaussians from the Probability course can treat those two proofs as optional and go straight to the worked example.

In one dimension, standardization and its inverse are

$$\begin{aligned} Z &= \frac{X - \mu}{\sigma} \sim \mathcal N(0, 1) \quad \text{for } X \sim \mathcal N(\mu, \sigma^2),\\[4pt] X &= \mu + \sigma Z \sim \mathcal N(\mu, \sigma^2) \quad \text{for } Z \sim \mathcal N(0, 1). \end{aligned}$$

So to sample any normal distribution it is enough to sample the standard one, stretch by $\sigma$ and shift by $\mu.$

Figure

Standardization and its inverse in one dimension

The bivariate case needs a matrix that plays the role of $\sigma$.

Spectral decomposition. By the spectral theorem, every real symmetric matrix, and in particular every covariance matrix, can be written as

$$\boldsymbol\Sigma = \mathbf U\,\boldsymbol\Lambda\,\mathbf U^{\top},$$

where $\boldsymbol\Lambda$ is diagonal with the eigenvalues of $\boldsymbol\Sigma$ and the columns of the orthogonal matrix $\mathbf U$ are the matching normalized eigenvectors. The eigenvalues are non-negative because $\boldsymbol\Sigma$ is positive semidefinite, so we can define the matrix square root

$$\boldsymbol\Sigma^{1/2} = \mathbf U\,\boldsymbol\Lambda^{1/2}\,\mathbf U^{\top},$$

where $\boldsymbol\Lambda^{1/2}$ holds the square roots of the eigenvalues. The root is itself symmetric, because $(\mathbf U\boldsymbol\Lambda^{1/2}\mathbf U^{\top})^{\top} = \mathbf U\boldsymbol\Lambda^{1/2}\mathbf U^{\top}$. Since $\mathbf U^{\top}\mathbf U = \mathbf I$, we get $\boldsymbol\Sigma^{1/2}\boldsymbol\Sigma^{1/2} = \mathbf U\boldsymbol\Lambda\mathbf U^{\top} = \boldsymbol\Sigma$.

Transformation. If $Z \sim \mathcal N(\mathbf 0, \mathbf I)$, then

$$X = \boldsymbol\mu + \boldsymbol\Sigma^{1/2} Z \sim \mathcal N(\boldsymbol\mu, \boldsymbol\Sigma).$$

More generally, any square matrix $\mathbf A$ with $\mathbf A\mathbf A^{\top} = \boldsymbol\Sigma$ works; $\boldsymbol\Sigma^{1/2}$ is one, because it is symmetric and squares to $\boldsymbol\Sigma$. For positive definite $\boldsymbol\Sigma$ such an $\mathbf A$ is invertible, since $(\det\mathbf A)^2 = |\boldsymbol\Sigma| > 0$, and the change-of-variables formula gives the density of $X = \boldsymbol\mu + \mathbf A Z$ (Murphy, 2022, §2.8). We substitute $\mathbf z = \mathbf A^{-1}(\mathbf x - \boldsymbol\mu)$ into the standard density $\frac{1}{2\pi}e^{-\mathbf z^{\top}\mathbf z/2}$ and divide by $|\det\mathbf A| = |\boldsymbol\Sigma|^{1/2}$:

$$\begin{aligned} p(\mathbf x) &= \frac{1}{2\pi\,|\det\mathbf A|}\exp\!\left(-\tfrac12\,\mathbf z^{\top}\mathbf z\right) \\[4pt] &= \frac{1}{2\pi\sqrt{|\boldsymbol\Sigma|}} \\[4pt] &\times \exp\!\left(-\tfrac12\, (\mathbf{x}-\boldsymbol\mu)^{\top}\,\boldsymbol\Sigma^{-1}\,(\mathbf{x}-\boldsymbol\mu)\right), \end{aligned}$$

because $\mathbf z^{\top}\mathbf z = (\mathbf x-\boldsymbol\mu)^{\top}(\mathbf A^{-1})^{\top}\mathbf A^{-1}(\mathbf x-\boldsymbol\mu)$ and $(\mathbf A^{-1})^{\top}\mathbf A^{-1} = (\mathbf A\mathbf A^{\top})^{-1} = \boldsymbol\Sigma^{-1}$. This is the density of $\mathcal N(\boldsymbol\mu, \boldsymbol\Sigma)$ from the previous section. (For a singular $\boldsymbol\Sigma$ the same map defines a degenerate Gaussian concentrated on a line, which has no density on the plane.) The first two moments give a quick check. The mean is $\boldsymbol\mu + \mathbf A\,\mathbb E[Z] = \boldsymbol\mu$, and the covariance is

$$\begin{aligned} \operatorname{Cov}(\boldsymbol\mu + \mathbf A Z) &= \mathbf A\,\operatorname{Cov}(Z)\,\mathbf A^{\top} \\ &= \mathbf A\,\mathbf I\,\mathbf A^{\top} = \boldsymbol\Sigma . \end{aligned}$$

The symmetric root $\boldsymbol\Sigma^{1/2}$ is one choice. For positive definite $\boldsymbol\Sigma$ there is exactly one lower-triangular $\mathbf L$ with positive diagonal and $\mathbf L\mathbf L^{\top} = \boldsymbol\Sigma$, the Cholesky factor (Golub and Van Loan, 2013, §4.2). It is the usual choice in code: a finite algorithm computes it in about $d^3/3$ floating-point operations for a $d \times d$ matrix, while an eigendecomposition needs an iterative method. Unlike the symmetric root, it requires positive definiteness: torch.linalg.cholesky fails on a singular covariance, and MultivariateNormal accepts only a positive definite covariance_matrix. PyTorch’s MultivariateNormal samples exactly this way: it stores the Cholesky factor of the covariance it receives and returns $\boldsymbol\mu + \mathbf L Z$ with $Z \sim \mathcal N(\mathbf 0, \mathbf I)$ (PyTorch documentation and source). The same pattern, noise in and a deterministic transformation out, is the reparameterization trick of the VAE and the one-step noising formula of DDPM.

Figure

Sampling any Gaussian from standard noise

Worked example#

For $\boldsymbol\Sigma = \begin{pmatrix} 1 & 0.5 \\ 0.5 & 1\end{pmatrix}$ the characteristic equation is $(1 - \lambda)^2 = 0.25$, so the eigenvalues are $\lambda_1 = 1.5$ and $\lambda_2 = 0.5$. The eigenvectors are $\tfrac{1}{\sqrt2}(1, 1)^{\top}$ and $\tfrac{1}{\sqrt2}(1, -1)^{\top}$. The variance is largest along the diagonal $x_1 = x_2$: the $1\sigma$ ellipse has semi-axes $\sqrt{1.5} \approx 1.225$ along it and $\sqrt{0.5} \approx 0.707$ across it. Two valid square roots are

$$\begin{aligned} \boldsymbol\Sigma^{1/2} &= \frac12\begin{pmatrix} \sqrt{1.5}+\sqrt{0.5} & \sqrt{1.5}-\sqrt{0.5} \\ \sqrt{1.5}-\sqrt{0.5} & \sqrt{1.5}+\sqrt{0.5} \end{pmatrix} \\[4pt] &\approx \begin{pmatrix} 0.966 & 0.259 \\ 0.259 & 0.966 \end{pmatrix}, \end{aligned}$$ $$\mathbf L = \begin{pmatrix} 1 & 0 \\ 0.5 & \sqrt{0.75} \end{pmatrix} \approx \begin{pmatrix} 1 & 0 \\ 0.5 & 0.866 \end{pmatrix}.$$

Check: $0.966^2 + 0.259^2 \approx 1$ and $2 \cdot 0.966 \cdot 0.259 \approx 0.5$. For the Cholesky factor, $\mathbf L\mathbf L^{\top} = \begin{pmatrix} 1 & 0.5 \\ 0.5 & 0.25 + 0.75 \end{pmatrix} = \boldsymbol\Sigma$. Both matrices turn independent standard noise into samples of the same correlated cloud.

Check yourself. Take $\boldsymbol\Sigma = \begin{pmatrix} 1 & -0.8 \\ -0.8 & 1\end{pmatrix}$. Along which diagonal does the cloud stretch, and what is its Cholesky factor? Answer: the eigenvalues are $1.8$ and $0.2$, the long axis runs along $x_1 = -x_2$ with semi-axes $\sqrt{1.8} \approx 1.342$ and $\sqrt{0.2} \approx 0.447$, and $\mathbf L = \begin{pmatrix} 1 & 0 \\ -0.8 & 0.6 \end{pmatrix}$, since $\mathbf L\mathbf L^{\top} = \begin{pmatrix} 1 & -0.8 \\ -0.8 & 0.64 + 0.36 \end{pmatrix} = \boldsymbol\Sigma$.

Gaussians in PyTorch#

torch.distributions contains parametrizable distributions with methods to evaluate their (log-)density and to draw samples. The 2-D Gaussian is MultivariateNormal. The fitting experiments in Fitting distributions with KL need a Gaussian whose parameters an optimizer can change freely. The covariance makes that non-trivial: an arbitrary $2\times 2$ matrix is generally not a valid covariance.

The fix is to learn the Cholesky factor instead of $\boldsymbol\Sigma$. With

$$\begin{gathered} \mathbf L = \begin{pmatrix} e^{a} & 0 \\ b & e^{c} \end{pmatrix}, \\[4pt] \mathbf L\mathbf L^{\top} = \begin{pmatrix} e^{2a} & b\,e^{a} \\ b\,e^{a} & b^2 + e^{2c} \end{pmatrix}, \end{gathered}$$

the three raw numbers $a, b, c$ can take any real values. The diagonal of $\mathbf L$ is stored through a logarithm, so it stays positive and the triangular matrix $\mathbf L$ is invertible. Then $\mathbf L\mathbf L^{\top}$ is symmetric and positive definite: for $\mathbf v \ne \mathbf 0$, $\mathbf v^{\top}\mathbf L\mathbf L^{\top}\mathbf v = \|\mathbf L^{\top}\mathbf v\|^2 > 0$ because $\mathbf L^{\top}$ is invertible. The running example is reached with $a = 0$, $b = 0.5$ and $c = \ln\sqrt{0.75} \approx -0.144$; the initialization $a = b = c = 0$ gives $\boldsymbol\Sigma = \mathbf I$.

Parameters that should be learned are created with requires_grad=True, which tells PyTorch’s autograd to record operations on them and compute their gradients. The model distribution $Q$ has a trainable mean and trainable raw entries of $\mathbf L$. The target distribution $P$ uses requires_grad=False, because it stays fixed. Figure 4 shows the starting point of the first fitting experiment, with each cloud coloured by its own density.

Figure

Target P and model Q before fitting

A minimal sketch of both ways to sample the running example:

import torch
from torch import distributions as distrs

mu = torch.zeros(2)
cov = torch.tensor([[1.0, 0.5], [0.5, 1.0]])
target = distrs.MultivariateNormal(
    mu, covariance_matrix=cov)

# 1000 samples, shape (1000, 2)
x = target.sample(torch.Size([1000]))
# L = [[1, 0], [0.5, 0.866]]
L = torch.linalg.cholesky(cov)
# the same distribution, by hand
x_manual = mu + torch.randn(1000, 2) @ L.T

Each row of torch.randn(1000, 2) is one draw $\mathbf z^{\top}$, so @ L.T computes $(\mathbf L\mathbf z)^{\top}$ row by row.

Notebook#

The practical part of this note follows the Distances between distributions notebook from the HSE-LAMBDA Deep Generative Models course repository. The link is pinned to a fixed commit, and cell numbers count every cell from 1.

• Cell 2 sets up the environment: it imports Matplotlib with the fivethirtyeight style, NumPy, PyTorch with torch.distributions as distrs, and IPython’s clear_output, which the training loops later use to redraw their plots.
• Cell 4 defines plot_2d_dots, a scatter plot of an $N \times 2$ array on the default limits $(-10, 10)$ for both axes (arguments xlim and ylim), one small dot per sample, coloured by an optional value array through the colour map cmap (default viridis).
• Cell 5 defines the three helpers behind the parametrization above. get_parameters(mu, sigma, requires_grad) creates the mean $(\mu, \mu)$ and the raw entries $a = c = \ln\sigma$, $b = 0$. create_distr assembles $\mathbf L$ and returns MultivariateNormal(mu, L @ L.T), which passes $\mathbf L\mathbf L^{\top}$ as covariance_matrix. sample(distr, n) returns an $n \times 2$ tensor. Since $\mathbf L$ is already a valid Cholesky factor, MultivariateNormal(mu, scale_tril=L) defines the same distribution and skips the re-factorization that PyTorch performs on L @ L.T; the PyTorch documentation recommends scale_tril for this reason.
• Cells 8–11 create the fixed target $\mathcal N((-5,-5), \mathbf I)$ and the trainable model $Q$ initialized at $\mathcal N(\mathbf 0, \mathbf I)$, and plot both clouds coloured by $p(\mathbf x)$ and $q(\mathbf x)$ (Figure 4). Cells 8 and 11 draw 1000 samples from each distribution, and the plots use the viridis and plasma colour maps; Figure 4 shows 400 samples per cloud and encodes density by colour strength instead. The fit itself is the subject of Fitting distributions with KL.

Next#

With a target distribution and a trainable model in hand, the remaining question is what “fit” means: we need a number that says how far $Q$ is from $P$. The first candidate is the KL divergence, and to understand it we start from entropy in Entropy, cross-entropy and KL divergence.

References#

• C. M. Bishop. Pattern Recognition and Machine Learning. Springer, 2006. §2.3: the Gaussian distribution and its moments.
• K. P. Murphy. Probabilistic Machine Learning: An Introduction. MIT Press, 2022. §2.8: transformations of random variables, including the change-of-variables formula; §3.2: the multivariate Gaussian and the Mahalanobis distance.
• G. H. Golub, C. F. Van Loan. Matrix Computations, 4th ed. Johns Hopkins University Press, 2013. §4.2: positive definite systems and the Cholesky factorization.
• PyTorch documentation: torch.distributions.MultivariateNormal; source of its sampling method, multivariate_normal.py (PyTorch 2.5.0).
• HSE-LAMBDA, Deep Generative Models course repository: distances.ipynb.