A simple 2d case of Gaussian simulation using Langevin Monte Carlo algorithm

Here’s the set up of multivariate Gaussian for a simple 2d example.

Let X be a bivariate Gassian random variable.

Let the mean of X $\mathbb{E}[X]=(0,0)$.

Let the covariance matrix of X be $$cov(X)= \left[ {\begin{array}{cc} 1 & 0.9 \\ 0.9 & 1 \\ \end{array} } \right]$$

Then the log density function of X is proportional to $-0.5*(X-\mathbb{E}[X])^T cov(X)^{-1} (X-\mathbb{E}[X])$. Note that $(X-\mathbb{E}[X])^T cov(X)^{-1} (X-\mathbb{E}[X])$ denotes the Mahalanobis distance between $X$ and its mean $\mathbb{E}[X]$. For multivariate normal distribution log density function derivation, see here

library(MASS)
require(ggplot2)

## Loading required package: ggplot2

library(mvtnorm)

# Define the target distribution: multivariate Gaussian
target_mean <- c(0, 0)
target_cov <- matrix(c(1, 0.9, 0.9, 1), nrow = 2) # covar matrix of the 2d Gaussian
target_log_density <- function(x) {
  -0.5 * t(x - target_mean) %*% solve(target_cov) %*% (x - target_mean)
}

Gradient of the log-density function is $cov(X)^{-1}(\mathbb{E}[X]-X)$

# Gradient of the log-density function
grad_log_density <- function(x) {
  solve(target_cov) %*% (target_mean - x)
}

Now make use of the Langevin Monte Carlo iteration, which states $X_{k+1} = X_k + \eta * \nabla(log(pdf(X))) / 2 + \sqrt{\eta}*Z_k$ where $Z_k$ is a random drawing from a bivariate normal distribution.

Question: if we need to make random draw from normal distribution for this LMC algorithm to work for normal distribution, doesn’t that mean it’s circulatory?

Answer:

Most of the time, LMC is used to sample random variables from a much more complicated distribution than Gaussian. So the step where we introduce the Gaussian noise to purturb the sequence of variables is not circulatory in those settings.

Question: what does “thinning” mean in the context of LMC sampling technique?

Answer:

Thinning refers to sampling random variables for a distribution by leaving several iterations in between unsampled. It is worth noting that

Question: How does strong convexity come into place when we are using LMC to sample a Gaussian distribution?

Answer:

For a function to be convex, the necessary condition is to check that the Hessian (or “second derivative” for 1d random variable) is positive semi-definite (or $\geq 0$ for 1d random variable).

Definition: A function $f: \Omega \subset \mathbb{R}^n \rightarrow \mathbb{R}$ is a Convex function iff $f(tx + (1-t)y) \leq tf(x) + (1-t)f(y)$.

(Second Order Condition for Convexity) For a function $f \in C^2(\Omega)$, where $\Omega \subset \mathbb{R}^n$ is also a convex set, $f$ is convex iff $\nabla^2f(x) \succeq 0$ for all $x \in \Omega$.

Definition: Let $f\in C^2(\Omega)$ and $f$ is convex. $f$ is strictly convex iff $\nabla^2 f(x) \succ 0$ for all $x \in \Omega$.

Definition: Let $f\in C^2(\Omega)$ and $f$ is convex. $f$ is strongly convex if there exists some $m > 0$ such that $\nabla^2 f(x) \succ mI$ for all $x \in \Omega$, where $I$ is just the identity matrix of the dimensions of Hessian $\nabla^2f(x)$.

So intuitively, the larger the second derivative/Hessian is, the more strongly convex a function is. Therefore, let’s calculate the second derivative of log of pdf of a standard normal random variable.

Let $x \sim N(0,1)$.

We know that log(x) is a concave function, so -log(x) is a convex function. The pdf of a standard normal is a convex function. So the composition of the two is also a convex function. Now let’s calculate how convex it is through second derivative calculation.

$$\begin{align*} -log(pdf(x)) &= -log(\frac{1}{\sqrt{2\pi}\sigma^2}exp(-\frac{(x-\mu)^2}{2\sigma^2}))\\ &= -log(\frac{1}{\sqrt{2\pi}\sigma^2}exp(-\frac{x^2-2x\mu+\mu^2}{2\sigma^2}))\\ &\propto -\frac{x^2-2x\mu+\mu^2}{2\sigma^2}\\ \implies \frac{d^2(-log(pdf(x)))}{dx^2} &\propto \frac{1}{\sigma^2} \end{align*}$$

This shows that the second derivative or Hessian is proportional to $\frac{1}{\sigma^2}$, so the smaller the variance $\sigma^2$, the bigger the second derivative, so the more strongly convex $-log(pdf(x))$ is.

We can relate this to the variance of $X_t$, as we want to achieve the convergence of $var(\lim_{t \rightarrow \infty} X_t) = 1$ since the simulated $X_t$ converges in distribution to $N(0,1)$.

Here’s how we can look at the variance in each iteration of the Langevin Monte Carlo algorithm:

$$\begin{align*} dX_t &= -\nabla(-log(pdf(X_t)))dt + \sqrt{2}dB_t\\ \implies X_{t+\epsilon}-X_t &= -\nabla(-log(pdf(X_t)))(t+\epsilon-t) + \sqrt{2} (B_{t+\epsilon}-B_{t})\\ \implies X_{t+\epsilon} &= X_t - \frac{X_t}{\sigma^2}\epsilon+\sqrt{2}(B_{t+\epsilon}-B_{t})\\ X_{t+\epsilon} &= X_t(1-\frac{\epsilon}{\sigma^2}) + \sqrt{2\epsilon}z_t \text{ where $z_t \sim N(0,1)$, and i.i.d. for each $t$}\\ \end{align*}$$

If we let the starting $X_0$ to be a random variable of $N(0,1)$, then we can tell the variance of $X_\epsilon$ by the above expression.

$$\begin{align*} X_\epsilon &= (1-\frac{\epsilon}{\sigma^2})X_0 + \sqrt{2\epsilon}Z)\\ &\sim N(0,(1-\frac{\epsilon}{\sigma^2})^2 + 2\epsilon) \end{align*}$$

So if we let $X_0 \sim N(0,var(0))$ where $var(0)$ denotes the variance. Then the variance at the $\epsilon$ step and at $t+\epsilon$ step can also be calculated.

$$\begin{align*} X_\epsilon &\sim N(0,(1-\frac{\epsilon}{\sigma^2})^2var(0) + 2\epsilon)\\ \implies var(\epsilon) &= (1-\frac{\epsilon}{\sigma^2})^2var(0) + 2\epsilon\\ X_{t+\epsilon} &\sim N(0,(1-\frac{\epsilon}{\sigma^2})^2var(t) + 2\epsilon))\\ \implies var(t+\epsilon) &= (1-\frac{\epsilon}{\sigma^2})^2var(t) + 2\epsilon \end{align*}$$

Now calculate the derivative of variance with respect to time to see how the variance changes overtime. We expect to express $\frac{dvar(t)}{dt}$ in terms of $var(t)$ so that we can solve the ODE to express $var(t)$ in terms of $t$.

$$\begin{align*} \lim_{\epsilon \rightarrow 0} \frac{var(t+\epsilon)-var(t)}{\epsilon} &= \lim_{\epsilon \rightarrow 0} \frac{(1-\frac{\epsilon}{\sigma^2})^2 var(t)+2\epsilon}{\epsilon}\\ &= \lim_{\epsilon \rightarrow 0} \frac{(-\frac{2\epsilon}{\sigma^2}+\frac{\epsilon^2}{\sigma^4})var(t)+2\epsilon}{\epsilon}\\ &= \lim_{\epsilon \rightarrow 0} (-\frac{2}{\sigma^2}+\frac{\epsilon}{\sigma^4})var(t)+2 \\ &= -\frac{2}{\sigma^2}var(t) + 2\\ \implies \frac{dvar(t)}{dt} &= -\frac{2}{\sigma^2}var(t) + 2 \end{align*}$$

Solving the ODE, we get $var(t) = \sigma^2(1-exp(-\frac{2}{\sigma^2}t))$.

And taking $t$ to infinity, we have $\lim_{t \rightarrow \infty} var(t) = \sigma^2$. Therefore, the smaller the $\sigma^2$, the smaller the variance is, which is what we would ideally want to converge to in infinite time. This is the same conclusion we drew from requiring the function $-log(pdf(x))$ to be strongly convex.

# Langevin dynamics sampling without "thinning"
langevin_sampling <- function(n_samples, step_size, initial_value) {
  samples <- matrix(0, n_samples, length(initial_value))
  x <- initial_value
  
  for (i in 1:n_samples) {
    grad <- grad_log_density(x)
    x <- x + step_size / 2 * grad + sqrt(step_size) * rnorm(length(x))
    samples[i, ] <- x
  }
  
  return(samples)
}

Next, generate examples using the langevin_sampling function defined above.

# Set parameters
n_samples <- 505000
bsz = 5000
step_size <- 0.01
initial_value <- c(0, 0)

# Generate samples
samples <- langevin_sampling(n_samples, step_size, initial_value)
samplesburn = samples[-c(1:bsz),]
thinsz=seq(100,n_samples-bsz,100)

colMeans(samples)

## [1] -0.04580787 -0.04383146

cov(samples)

##           [,1]      [,2]
## [1,] 0.9814081 0.8814491
## [2,] 0.8814491 0.9863900

samplesDf = data.frame(samples)

library(ggplot2)
ggplot(samplesDf, aes(x=X1,y=X2))+
  labs(title = "Langevin Sampling for 2 dimensions",
       x = "x1",
       y = "x2") +
  geom_point(color = "deepskyblue2", alpha = 0.5) + 
  theme_bw() 

References

Simsekli et al., Stochastic Quasi-Newton Langevin Monte Carlo (2016)

Nguyen et al. Non-Asymptotic Analysis of Fractional Langevin Monte Carlo for Non-Convex Optimization (2019)