variance of conditional multivariate gaussian

Question

I was playing around with Gaussian Distributions on my machine and I was interesting in making a pretty plot. I wanted to show the distribution of $x_1$ if $x_2$ was given if $x_1,x_2$ were distributed by a multivariate normal distribution.

The image feels "wrong". I would imagine that the red likelihood in the center of the blue distribution would be wider than at the edge. After double checking the code, I notice it might be the maths.

Maths on wikipedia as well as my books confirm that the distribution of $x_1$ conditional on $x_2$ = $a$ is multivariate normal $(x_1 | x_2 = a) \sim N(\hat{\mu}, \hat{\Sigma})$ where

$$ \bar{\boldsymbol\mu} = \boldsymbol\mu_1 + \boldsymbol\Sigma_{12} \boldsymbol\Sigma_{22}^{-1} \left( \mathbf{a} - \boldsymbol\mu_2 \right) $$

and covariance matrix

$$ \overline{\boldsymbol\Sigma} = \boldsymbol\Sigma_{11} - \boldsymbol\Sigma_{12} \boldsymbol\Sigma_{22}^{-1} \boldsymbol\Sigma_{21}. $$

When looking at the maths, it seems that the variance of $p(x_1|x_2 = a)$ does not depend on the value of $a$. This is starting to feel very counter intuitive so I am wondering if I am missing something.

Edit

To clarify the question: I am aware that the maths are correct and I can follow the mathematics but there is something counter-intuitive about the result that I just cannot get my head around. When I look at the 2d distribution, suppose that $x_2=a$ is given, I would expect the variance to change for different values of $a$.

From just starting at the 2d plot it seems wrong to consider that the variance of $p(x_1 | x_2 = a)$ is constant and independant of $a$. What am I missing here?

The code that generated the plot

import matplotlib.pylab as plt 

import torch
from torch.distributions import Normal as norm
from torch.distributions.multivariate_normal import MultivariateNormal as mvnorm

#@title different given values { run: "auto" }
g1 = -3.5 #@param {type:"slider", min:-4, max:4, step:0.1}
g2 = 0.2 #@param {type:"slider", min:-4, max:4, step:0.1}
g3 = 2.9 #@param {type:"slider", min:-4, max:4, step:0.1}

m = torch.tensor([0.0, 0.0])
c = torch.tensor([[1.0, 0.9], [0.9, 1.0]])

s = mvnorm(m, c).sample(sample_shape=(5000,))
s_np = s.numpy().reshape(5000, 2)
plt.figure(figsize=(6,5))
plt.scatter(s_np[:, 0], s_np[:, 1], alpha=0.3)

for g in [g1, g2, g3]:
    mu_pred = m[1] + c[0][1]/c[1][1]*(g - m[0])
    sigma_pred = c[1][1] - c[1][0]/c[0][0]*c[0][1]
    fitted_distr = norm(mu_pred, sigma_pred)
    print(f"g:{g:.3}, mu:{mu_pred:.2}, sigma:{sigma_pred:.4}")

    xs = torch.linspace(-4, 4, 300)
    likelihood = torch.exp(fitted_distr.log_prob(xs)).numpy()

    plt.plot(xs.numpy(), g + likelihood, c='red')

Answer 1

The problem is the sampling, which is to say that you're trying to compare the conditional density with the distribution of slices of data that were sampled from a bivariate density.

The theoretical variance of $x_1|x_2$ doesn't depend on the value of $x_2$ but the observed does because the further out you get from the means of $x_1$ and $x_2$, the fewer samples you will have. So what you're seeing is still 100% correct, but the quantity of points is creating the illusion of narrow and wide variances.

Using the density of $x_1|x_2$ from the same Wikipedia section you linked to:

$$ (X_1\mid X_2=x_2) \ \sim\ \mathcal{N}\!\left(\mu_1+\frac{\sigma_1}{\sigma_2}\rho( x_2 - \mu_2),\, (1-\rho^2)\sigma_1^2\right) $$

we can compare data sampled bivariate-ly with data sampled conditionally:

library(magrittr)
library(ggplot2)

# multivariate-sampled data:
mu <- c(0, 0)
S <- matrix(c(1, 0.9, 0.9, 1), nrow = 2, byrow = TRUE)
set.seed(0)
xs <- mvtnorm::rmvnorm(10000, mean = mu, sigma = S) %>%
  set_colnames(c("X1", "X2"))
df <- as.data.frame(xs)

# functions & data to overlay the conditional density:
conditional_dnorm <- function(x1, x2, m1, m2, s1, s2, r) {
  mean <- m1 + (s1/s2) * r * (x2 - m2)
  var <- (1 - r^2) * s1^2
  sd <- sqrt(var)
  return(dnorm(x1, mean, sd))
}
conditional_probabilities <- purrr::map_dfr(-2:2, function(x2) {
  x1 <- seq(-4, 4, length.out = 1000)
  fx1 <- purrr::map_dbl(
    x1, conditional_dnorm, x2 = x2,
    m1 = 0, m2 = 0, s1 = 1, s2 = 1, r = 0.9
  )
  return(data.frame(X1 = x1, X2 = x2, FX1 = fx1))
})

# plot:
ggplot(df, aes(x = X1, y = X2)) +
  geom_point(alpha = 0.1) +
  geom_line(
    aes(x = X1, y = FX1 + X2, group = X2),
    data = conditional_probabilities,
    color = "red"
  ) +
  theme_minimal() +
  labs(title = "Sampled multivariately")

That's just to confirm your results. The distribution of $x_1$ at various values of $x_2$ does look different. But, let's sample $x_1$ conditionally on $x_2$, ensuring that at each value of $x_2$ we get the same quantity of samples of $x_1$:

conditional_rnorm <- function(n, x2, m1, m2, s1, s2, r) {
  mean <- m1 + (s1/s2) * r * (x2 - m2)
  var <- (1 - r^2) * s1^2
  sd <- sqrt(var)
  return(rnorm(n, mean, sd))
}
conditional_samples <- purrr::map_dfr(
  seq(-2, 2, length.out = 100),
  function(x2) {
    x1 <- conditional_rnorm(100, x2, 0, 0, 1, 1, 0.9)
    return(data.frame(X1 = x1, X2 = x2))
  }
)

ggplot(conditional_samples, aes(x = X1, y = X2)) +
  geom_point(alpha = 0.1) +
  geom_line(
    aes(x = X1, y = FX1 + X2, group = X2),
    data = conditional_probabilities,
    color = "red"
  ) +
  theme_minimal() +
  labs(title = "Sampled conditionally")

Looks better, doesn't it? Hope that helps!

Answer 2

The situation is addressed in the following paragraphs of your wikipedia link. It is also easier to look at Bivariate case there, since it directly gives the formulas you need. The conditional variance is $(1-\rho^2)\sigma_1^2$.

Using the bivariate formula in here and assuming zero-means for simplicity, wlog, we can factorize the joint gaussian pdf to find the conditional gaussian pdf for $x|y$:

$$p(x,y)=p(x|y)p(y)\propto exp(-\frac{1}{2(1-\rho^2)}(\frac{x_1^2}{\sigma_1^2}-\frac{2\rho x_1x_2}{\sigma_1\sigma_2}+\frac{x^2_2}{\sigma^2_2}))$$ $$\propto exp(-\frac{1}{2(1-\rho^2)}(\frac{x_1^2}{\sigma_1^2}-\frac{2\rho x_1x_2}{\sigma_1\sigma_2}+\frac{x^2_2}{\sigma^2_2}-\frac{(1-\rho^2)x_2^2}{\sigma_2^2}))exp(-\frac{x_2^2}{2\sigma_2^2})$$

Second part is $p(y)$, first part is $p(x|y)$ and both are in still normal form. Normal form is proportional to $exp(-\frac{(x-\mu)^2}{2\sigma^2})=exp(-\frac{x^2}{2\sigma^2}+\frac{x\mu}{\sigma^2}-\frac{\mu^2}{\sigma^2})$. We just match the insides of the exponentials here. The square term in exponent of $p(x|y)$ is $-\frac{x_1^2}{2(1-\rho^2)\sigma_1^2}$, which should be equal to $-\frac{x_1^2}{2\sigma^2}$, which yields $\sigma^2=(1-\rho^2)\sigma_1^2$, which is the same formula above and the variance doesn't depend on the particular value of $x_2$. Notice that I assumed zero-mean, but incorporating $\mu_1$ and $\mu_2$ into these equations will not change the coefficient of $x_1^2$.