Today's Topics

• Importance sampling

• Measure transport sampler

Importance sampling

Strictly speaking, importance sampling (IS) is not a method to obtain sample points $X_1,\dots ,X_n$ that follow the target distribution $p\left(x\right)$.

Instead, it is a technique to estimate expectations related to $p\left(x\right)$:

$$\mu =E_X\left[f\right(X\left)\right]=\int f\left(x\right)p\left(x\right)dx,X\sim p\left(x\right).$$

A Direct Solution

Of course, one direct method to approximate $\mu$ is to generate $X_1,\dots ,X_M\sim p\left(x\right)$, and then an unbiased estimator for $\mu$ is given by

$$\hat{\mu }=\frac{1}{M}\sum _{i=1}^{M}f\left(X_i\right),X_i\sim p\left(x\right),i=1,\dots ,M.$$

Suppose that we use rejection sampling to get $X_i$ based on a proposal distribution $q\left(x\right)$. There are two issues here:

• Rejection sampling discards sample points
• $f\left(x\right)$ may be close to zero outside a region $A$ for which $P(X\in A)$ is small

Example

We want to estimate $p=P(X>\pi )$ and $\mu =E(X|X>\pi )$, where $X\sim N(0,1)$.

Naive solution: sample $X_1,\dots ,X_M\stackrel{iid}{\sim }N(0,1)$, and get $$\begin{array}{cc}\hat{p}& =\frac{1}{M}\sum _{i=1}^{M}I\{X_i>\pi \},\\ \hat{\mu }& =\frac{{\sum }_{i=1}^M{X}_i\cdot I\{X_i>\pi \}}{{\sum }_{i=1}^{M}I\{X_i>\pi \}}.\end{array}$$

Problem: $p$ is very small (true value ~0.00084), so maybe all $X_i$'s are smaller than $\pi$.

Example

est_naive = function(n)
{
    x = rnorm(n)
    p_hat = mean(x > pi)
    mu_hat = sum(x * (x > pi)) / sum(x > pi)
    c(p_hat, mu_hat)
}
set.seed(123)
est_naive(n = 100)

## [1]   0 NaN

Motivation

IS attempts to resolve the previous issues:

• It does not discard or waste any sample points

• Instead, it assigns different weights to each point

• By properly choosing the proposal distribution, it is able to more effectively generate sample points around the "important region" $A$

Basic Idea

The idea of IS is in fact quite simple. It is based on a straightforward identity:

$$\mu =\int f\left(x\right)p\left(x\right)dx=\int \frac{f\left(x\right)p\left(x\right)}{q\left(x\right)}q\left(x\right)dx=E_q\left(\frac{f\left(X\right)p\left(X\right)}{q\left(X\right)}\right),$$

where $q\left(x\right)$ is another density function that is positive on $R^p$, and $E_q(\cdot )$ denotes the expectation for $X\sim q\left(x\right)$.

Accordingly, the IS estimate for $\mu$ is

$${\hat{\mu }}_q=\frac{1}{M}\sum _{i=1}^M\frac{f\left(X_i\right)p\left(X_i\right)}{q\left(X_i\right)},X_i\sim q\left(x\right),i=1,\dots ,M.$$

Theorem ^[1]

Suppose that $q\left(x\right)>0$ whenever $f\left(x\right)p\left(x\right)\ne 0$. Then $E_q\left({\hat{\mu }}_q\right)=\mu$, and ${Var}_q\left({\hat{\mu }}_q\right)={\sigma }_q^2/n$, where

$${\sigma }_q^2={\int }_Q\frac{\left[f\right(x\left)p\right(x){]}^2}{q\left(x\right)}dx-{\mu }^2={\int }_Q\frac{\left[f\right(x\left)p\right(x)-\mu q(x){]}^2}{q\left(x\right)}dx,$$

and $Q=\{x:q(x)>0\}$.

Example

Back to the previous example, note that

$$\begin{array}{cc}p& ={\int }_{-\infty }^{+\infty }I(x>\pi )\varphi \left(x\right)dx\\ & ={\int }_{\pi }^{+\infty }\frac{I(x>\pi )\varphi \left(x\right)}{q\left(x\right)}q\left(x\right)dx={\int }_{\pi }^{+\infty }\frac{\varphi \left(x\right)}{q\left(x\right)}q\left(x\right)dx,\\ q\left(x\right)& =\{\begin{array}{cc}0,& x\le \pi \\ e^{-(x-\pi )},& x>\pi \end{array},\end{array}$$

where $q\left(x\right)$ is a shifted exponential distribution. Similarly,

$$\mu =p^{-1}{\int }_{\pi }^{+\infty }\frac{x\cdot \varphi \left(x\right)}{q\left(x\right)}q\left(x\right)dx.$$

Example

est_is = function(n)
{
    x = rexp(n) + pi
    ratio = exp(dnorm(x, log = TRUE) + x - pi)
    p_hat = mean(ratio)
    mu_hat = mean(x * ratio) / p_hat
    c(p_hat, mu_hat)
}
set.seed(123)
est_is(n = 100)

## [1] 0.0007478989 3.4296517837

Optimal $q\left(x\right)$ ^[1]

It can be proved that the optimal proposal distribution $q^{\ast }\left(x\right)$ is given by $q^{\ast }\left(x\right)=|f\left(x\right)|p\left(x\right)/E_p\left(|f\right(X\left)|\right)$.

Proof: for any density function $q\left(x\right)$ such that $q\left(x\right)>0$ when $f\left(x\right)p\left(x\right)\ne 0$,

$$\begin{array}{cc}{\mu }^2+{\sigma }_{q^{\ast }}^2& =\int \frac{\left[f\right(x\left)p\right(x){]}^2}{q^{\ast }\left(x\right)}dx=\left[E_p\right(|f\left(X\right)|){]}^2\\ & ={\left[E_q\left(\frac{|f\left(X\right)|p\left(X\right)}{q\left(X\right)}\right)\right]}^2\\ & \le E_q\left(\frac{\left[f\right(X\left)p\right(X){]}^2}{\left[q\right(X){]}^2}\right)=\int \frac{\left[f\right(x\left)p\right(x){]}^2}{q\left(x\right)}dx={\mu }^2+{\sigma }_q^2.\end{array}$$

Optimal $q\left(x\right)$

This means that to approximate $\int f\left(x\right)p\left(x\right)dx$, IS can be better than the simple Monte Carlo estimator!

Self-Normalized IS

Suppose that we can only compute $p_u\left(x\right)\propto p\left(x\right)$ and $q_u\left(x\right)\propto q\left(x\right)$, the self-normalized IS estimate is given by

$$\tilde{\mu }=\frac{{\sum }_{i=1}^{M}f\left(X_i\right)w\left(X_i\right)}{{\sum }_{i=1}^{M}w\left(X_i\right)},$$

where $w\left(x\right)=p_u\left(x\right)/q_u\left(x\right)$ and $X_i\sim q\left(x\right)$.

Under mild conditions, $\tilde{\mu }$ is a consistent estimator of $\mu$, but in general $\tilde{\mu }$ is no longer unbiased.

Recap: Inverse Transform Algorithm

• $U\sim Unif(0,1)$

• $g=F^{-1}$

• $X=g\left(U\right)\Rightarrow X\sim f\left(x\right)$

• $f\left(x\right)=F^{\prime }\left(x\right)$ is the density function

Random Variable Transformation

• Continuous random variable $X\sim p_X\left(x\right)$

• $p_X\left(x\right)$ density function

• $g:R\to R$ a monotone function

• Define $Y=g\left(X\right)$, then its density function is given by $$p_Y\left(y\right)=p_X\left(g^{-1}\right(y\left)\right)\left|\frac{d}{dy}{g}^{-1}\right(y\left)\right|$$

• Can extend to multivariate case

Random Vector Transformation

• Continuous random vector $X\in R^d$, $X\sim p_X\left(x\right)$

• $p_X\left(x\right)$ density function

• $T:R^d\to R^d$ a diffeomorphism (a smooth mapping with smooth inverse)

• Define $Y=T\left(X\right)$, then its density function is given by $$\begin{array}{cc}{p}_Y\left(y\right)& =p_X\left(T^{-1}\right(y\left)\right)|det\left(\nabla \right(T^{-1}\left)\right(y\left)\right)|\\ & =p_X\left(x\right){|det\left(\nabla T\right(x\left)\right)|}^{-1},x=T^{-1}\left(y\right)\end{array}$$

• $\nabla T$ (cf. $\nabla T^{-1}$) is the Jacobian matrix of $T$ (cf. $T^{-1}$)

Recall: Box-Muller Transform

1. $(U_1,U_2)\sim Unif\left(\right[0,1{]}^2)$

2. Let $Z_1=\sqrt{-2\log \left(U_1\right)}\cos \left(2\pi U_2\right)$ and $Z_2=\sqrt{-2\log \left(U_1\right)}\sin \left(2\pi U_2\right)$

3. Then $Z_1$ and $Z_2$ are two independent $N(0,1)$ random variables

Recall: Box-Muller Transform

• Transformation mapping $$T\left(\begin{array}{c}{U}_1\\ U_2\end{array}\right)=\left(\begin{array}{c}\sqrt{-2\log \left(U_1\right)}\cos \left(2\pi U_2\right)\\ \sqrt{-2\log \left(U_1\right)}\sin \left(2\pi U_2\right)\end{array}\right)$$

• Jacobian matrix $$\begin{array}{cc}\nabla T\left(\begin{array}{c}{U}_1\\ U_2\end{array}\right)& =\left(\begin{array}{cc}\frac{-\cos \left(\theta \right)}{U_{1}L}& -2\pi L\sin \left(\theta \right)\\ \frac{-\sin \left(\theta \right)}{U_{1}L}& 2\pi L\cos \left(\theta \right)\end{array}\right)\\ L& =\sqrt{-2\log \left(U_1\right)},\theta =2\pi U_2\end{array}$$

• Determinant $$det\nabla T=-2\pi U_1^{-1}$$

Recall: Box-Muller Transform

• Clearly, $p_U(u_1,u_2)=1$, so $$p_Z(z_1,z_2)=p_U(u_1,u_2)|det\nabla T{|}^{-1}=(2\pi {)}^{-1}{u}_1$$

• To express $u_1$ using $z_1$ and $z_2$, note that $$Z_1^2+Z_2^2=-2\log \left(U_1\right)$$

• Therefore, $$p_Z(z_1,z_2)=\frac{1}{2\pi }{e}^{-\frac{1}{2}(z_1^2+z_2^2)},$$ which implies that $Z_1$ and $Z_2$ follow independent $N(0,1)$

Ideas from Inverse Transform

• Fix a source distribution $p_Z\left(z\right)$, e.g. $N(0,I)$

• Given a target distribution $p\left(x\right)$

• Can we learn a mapping $T$ such that $X=T\left(Z\right)\sim p\left(x\right)$ if $Z\sim p_Z\left(z\right)$?

Measure Transport Sampler

• If we can obtain such a mapping, sampling would be easy:

  • Simulate $Z_1,\dots ,Z_n\sim N(0,I)$

  • Set $X_1=T\left(Z_1\right),\dots ,X_n=T\left(Z_n\right)$

  • Then $X_i\stackrel{iid}{\sim }p\left(x\right)$

• The key is to find such a mapping $T$, also called a transport map

Estimating Transport Map

• Fix a source distribution $Z\sim p_0\left(z\right)$, e.g. $N(0,I)$

• Assume we have a model $T_{\theta }\in T$ to parameterize $T$

• $T$ is chosen such that $T_{\theta }$ is a diffeomorphism (we will cover this later)

• Then $X=T_{\theta }\left(Z\right)\sim p_{\theta }\left(x\right)$ has a known density function

• The problem becomes how to estimate $\theta$

Criterion: KL Divergence

• Given target distribution $p\left(x\right)$

• We want to make $p_{\theta }\left(x\right)$ close to $p\left(x\right)$ as much as possible

• A natural criterion is the Kullback-Leibler (KL) divergence $$D_{KL}(p_{\theta }\parallel p)=\int p_{\theta }\left(x\right)\log \frac{p_{\theta }\left(x\right)}{p\left(x\right)}dx=E_{p_{\theta }}\left(\log \frac{p_{\theta }}{p}\right)$$

• The best $\theta$ is given by $${\theta }^{\ast }=\underset{\theta }{\arg min}{D}_{KL}(p_{\theta }\parallel p)$$

Training Algorithm

• An unbiased estimator for $D_{KL}(p_{\theta }\parallel p)$ is given by $$\hat{L}\left(\theta \right)=\frac{1}{M}\sum _{i=1}^M\left[\log p_{\theta }\right(X_i)-\log p(X_i\left)\right],$$ where $X_i=T_{\theta }\left(Z_i\right),Z_i\sim N(0,I)$

• We further have $$p_{\theta }\left(x\right)=p_0\left(z\right){|det\left(\nabla T_{\theta }\right(z\left)\right)|}^{-1},z=T_{\theta }^{-1}\left(x\right)$$

Training Algorithm

• As a result, $$\begin{array}{cc}\hat{L}\left(\theta \right)& =\frac{1}{M}\sum _{i=1}^M\left[\log p_0\right(Z_i)-\log |det\left(\nabla T_{\theta }\right(Z_i\left)\right)|-\log p(X_i\left)\right]\\ & =-\frac{1}{M}\sum _{i=1}^M[\log |det\left(\nabla T_{\theta }\right(Z_i\left)\right)|+\log p(T_{\theta }\left(Z_i\right)\left)\right]+C\end{array}$$

• Therefore, $$\begin{array}{cc}{E}_Z\hat{L}\left(\theta \right)& =D_{KL}(p_{\theta }\parallel p)\\ E_Z\left({\nabla }_{\theta }\hat{L}\right(\theta \left)\right)& ={\nabla }_{\theta }{D}_{KL}(p_{\theta }\parallel p)\end{array}$$

Stochastic Gradient Descent

We can then use stochastic gradient descent to optimize $D_{KL}(p_{\theta }\parallel p)$.

In each iteration:

• Simulate $Z_i\stackrel{iid}{\sim }N(0,I)$

• Define $$\hat{\ell }\left(\theta \right)=-\frac{1}{M}\sum _{i=1}^M[\log |det\left(\nabla T_{\theta }\right(Z_i\left)\right)|+\log p(T_{\theta }\left(Z_i\right)\left)\right]$$

• Compute $g\left(\theta \right)={\nabla }_{\theta }\hat{\ell }\left(\theta \right)$

• Update $\theta \leftarrow \theta -\alpha \cdot g\left(\theta \right)$

Unnormalized Target Density

• In fact, the training algorithm is also valid for unnormalized target densities.

• Assume $p\left(x\right)\propto e^{-E\left(x\right)}$

• Then just change $\hat{\ell }\left(\theta \right)$ to $$\hat{\ell }\left(\theta \right)=-\frac{1}{M}\sum _{i=1}^M[\log |det\left(\nabla T_{\theta }\right(Z_i\left)\right)|-E(T_{\theta }\left(Z_i\right)\left)\right]$$

• This is because the unknown constant does not affect the gradient

Demonstration

Modeling Transport Maps

For practical use, the transport map $T_{\theta }$ needs to have some "nice" properties:

• $T_{\theta }$ should be invertible and differentiable

• $T_{\theta }^{-1}$ and $det\left(\nabla T_{\theta }\right)$ should be easy to compute

• $T_{\theta }$ should be flexible enough to characterize sophisticated nonlinear mappings

At first glance, these are quite strong conditions.

Modeling Transport Maps

• Polynomials ^[2] (not good enough)

• Normalizing flows ^[3] (tools from the deep learning community)

Normalizing Flow

• Normalizing flows (NFs) are a class of diffeomorphisms constructed by neural networks

• NFs define invertible mappings through composition: $$T=T_K\circ \cdots \circ T_1$$

• Each $T_i$ is a simpler diffeomorphism

Change-of-Variable Revisited

• Recall the density formula for $X=T\left(Z\right)$: $$p_X\left(x\right)=p_Z\left(z\right){|det\left(\nabla T\right(z\left)\right)|}^{-1},z=T^{-1}\left(x\right)$$
• For $T=T_K\circ \cdots \circ T_1$, let $z_0=z$, $z_K=x$, and $z_i=T_i\left(z_{i-1}\right)$, then $$\log |det\left(\nabla T\right(z\left)\right)|=\sum _{i=1}^K\log |det\left(\nabla T_i\right(z_{i-1}\left)\right)|$$
• So we still obtain an explicit form for $p_X\left(x\right)$

Implementations

Many different implementations of NFs ^[4,5]:

• Permutation and orthogonal

• Decomposition-based

• Planar and radial

• Coupling

• Autoregressive

• ...

Example: Real NVP

• One popular implementation of NF is the affine coupling flow, also called Real NVP ^[6]

• Given $X\in R^d$, define $Y=T\left(X\right)\in R^d$ as below: $$\begin{array}{cc}{Y}_{1:r}& =X_{1:r}\\ Y_{(r+1):d}& =\mu \left(X_{1:r}\right)+\sigma \left(X_{1:r}\right)\odot X_{(r+1):d}\end{array}$$

• $X_{1:r}$ is the first $r$ elements of $X$, $r<d$

• $\odot$ is elementwise multiplication

• $\mu ,\sigma :R^r\to R^{d-r}$ are two neural networks, $\sigma (\cdot )>0$

Why Real NVP Works

Consider the composition of two transformations: $$\begin{array}{cc}\left(\begin{array}{c}{X}_1\\ X_2\\ X_3\\ X_4\end{array}\right)& \stackrel{T_1}{\Rightarrow }\left(\begin{array}{c}{Y}_1=X_1\\ Y_2=X_2\\ Y_3={\mu }_{Y_3}(X_1,X_2)+{\sigma }_{Y_3}(X_1,X_2)\cdot X_3\\ Y_4={\mu }_{Y_4}(X_1,X_2)+{\sigma }_{Y_4}(X_1,X_2)\cdot X_4\end{array}\right)\\ \left(\begin{array}{c}{Y}_1\\ Y_2\\ Y_3\\ Y_4\end{array}\right)& \stackrel{T_2}{\Rightarrow }\left(\begin{array}{c}{Z}_1={\mu }_{Z_1}(Y_3,Y_4)+{\sigma }_{Z_1}(Y_3,Y_4)\cdot Y_1\\ Z_2={\mu }_{Z_2}(Y_3,Y_4)+{\sigma }_{Z_2}(Y_3,Y_4)\cdot Y_2\\ Z_3=Y_3\\ Z_4=Y_4\end{array}\right)\end{array}$$

Why Real NVP Works

$T_1$ can be easily inverted (similar for $T_2$): $$\begin{array}{c}\left(\begin{array}{c}{X}_1\\ X_2\\ X_3\\ X_4\end{array}\right)\stackrel{T_1}{\Rightarrow }\left(\begin{array}{c}{Y}_1=X_1\\ Y_2=X_2\\ Y_3={\mu }_{Y_3}(X_1,X_2)+{\sigma }_{Y_3}(X_1,X_2)\cdot X_3\\ Y_4={\mu }_{Y_4}(X_1,X_2)+{\sigma }_{Y_4}(X_1,X_2)\cdot X_4\end{array}\right)\\ \left(\begin{array}{c}{X}_1=Y_1\\ X_2=Y_2\\ X_3=[Y_3-{\mu }_{Y_3}(Y_1,Y_2\left)\right]/{\sigma }_{Y_3}(Y_1,Y_2)\\ X_4=[Y_4-{\mu }_{Y_4}(Y_1,Y_2\left)\right]/{\sigma }_{Y_4}(Y_1,Y_2)\end{array}\right)\stackrel{T_1^{-1}}{\Leftarrow }\left(\begin{array}{c}{Y}_1\\ Y_2\\ Y_3\\ Y_4\end{array}\right)\end{array}$$

Why Real NVP Works

$\nabla T_1$ is lower triangular: $$\begin{array}{cc}\left(\begin{array}{c}{X}_1\\ X_2\\ X_3\\ X_4\end{array}\right)& \stackrel{T_1}{\Rightarrow }\left(\begin{array}{c}{Y}_1=X_1\\ Y_2=X_2\\ Y_3={\mu }_{Y_3}(X_1,X_2)+{\sigma }_{Y_3}(X_1,X_2)\cdot X_3\\ Y_4={\mu }_{Y_4}(X_1,X_2)+{\sigma }_{Y_4}(X_1,X_2)\cdot X_4\end{array}\right)\\ \nabla T_1\left(x\right)& =\left(\begin{array}{cccc}1& 0& 0& 0\\ 0& 1& 0& 0\\ \ast & \ast & {\sigma }_{Y_3}(x_1,x_2)& 0\\ \ast & \ast & 0& {\sigma }_{Y_4}(x_1,x_2)\end{array}\right)\end{array}$$

$det\left(\nabla T_1\right)$ can be computed in linear time (the product of diagonal elements).

Why Real NVP Works

Some theoretical works such as [7] show that Real NVP flows are universal approximators.

As a result, all three conditions are satisfied.

Parameterization

• Construct the transport map $T_{\theta }$ using normalizing flows

• The vector $\theta$ contains all the neural network parameters

• ${\nabla }_{\theta }\hat{\ell }\left(\theta \right)$ is typically computed using automatic differentiation

Extensions

Recent literature such as [8] shows that the measure transport sampler based on the KL loss function does not learn multimodal distributions well.

Some improved samplers exist, and more is yet to be explored.

References

References