INLA

Approximating the full joint posterior distribution using a Gaussian distribution is inaccurate. An alternative is to approximate the marginal posterior distribution of some subset of the parameters, referred to as the marginal Laplace approximation.

Then, integrate out the remaining parameters using another method.

• Integrated. Using numerical integration
• Nested. Because we need $p(\theta |y)$ to get $p(u|y)$
• Laplace approximations. The methods used to obtain parameters for the Gaussian approximation.

When is it useful?

Wherever we have a latent Gaussian. Almost always.

Following Dan Simpson, we can write a linear mixed model as

$$y=X\beta +Zb+ϵ,$$

where

• Fixed effects. $X$ contain the covariates, $\beta$ is a vector of unknown regression coefficients with a Gaussian prior $\beta \sim N(0,R)$
• Random effects. $Z$ is a known matrix that describes the random effects, $b\sim N(0,{\Sigma }_b)$

We can put fixed and random effects together $A=[X⋮Z]$, with a latent field (or latent Gaussian component) $u=({\beta }^T,b^T{)}^T$.

This allows us to re-write the model as a three-stage hierarchical model

$$p(y,u,\theta )=p(y|u,\theta )p(u|\theta )p(\theta ),$$

where $\mathbf{u}=(u_1,\dots ,u_N)$ is the latent field, and $\theta =({\theta }_1,\dots ,{\theta }_m)$ are the hyperparameters. For the latent field, $p(u|\theta )$, we can write

$$u\mid \theta \sim N(0,Q(\theta {)}^{-1})$$

and parametrise by a precision matrix

$$Q(\theta )=\left(\begin{array}{cc}{\Sigma }_b^{-1}& 0\\ 0& R^{-1}\end{array}\right).$$

What am I approximating?

Marginalise out $u$ and then use standard inference techniques on $\theta$.

The posterior can be written as

$$p(\theta |y)=\frac{p(y|u,\theta )p(u|\theta )p(\theta )}{p(u|y,\theta )}\times \frac{1}{p(y)}$$

• GMRF specifies $p(u|\theta )$
• prior $p(\theta )$ is supplied
• measurement model gives us $p(y|u,\theta )$
• normalising constant $p(y)$ is ignored since $y$ is fixed

log_prior_fn(params) + log_marginal_likelihood(params, y)

The other term in the denominator, $p(u|y,\theta )$, is estimated with Laplace’s method, $p(u|y,\theta )\approx \mathcal{N}(u^{\ast },{\Sigma }^{\ast })$ where $u^{\ast }$ matches the mode and $[{\Sigma }^{\ast }{]}^{-1}$ the curvature

The likelihood

$$p(y|\theta )=\frac{p(y|u,\theta )p(u|\theta )}{p(u|y,\theta )}$$

# evaluated at x = \hat{x}_0 below
log_laplace_approx = self._conditional_gaussian_approximation(
	x, y, gaussian
).logpdf(x)
log_marginal_likelihood = (    # P(y | params) =
	gaussian.logpdf(x)         # P(x | params)
	+ self.f(x, y).sum()       # * P(y | x, params)
	- log_laplace_approx       # / P(x | y, params)
)

Laplace approximation

Denoting $f(x)=\log (p(y|u))$, we have the following approximation around any given $\hat{x}$:

$$\begin{array}{rl}\log (p(x|y,\theta ))& =\log (p(y|x,\theta ))+\log (p(x|\theta ))+\text{const}\\ & =f(x)-\frac{1}{2}(x-\mu {)}^{T}Q(x-\mu )+\frac{1}{2}\log \det (Q)+\text{const}\\ & \approx f(\hat{x})+(x-\hat{x})f^{\prime }(\hat{x})+\frac{1}{2}(x-\hat{x}{)}^2{f}^{\prime \prime }(\hat{x})-\frac{1}{2}(x-\mu {)}^{T}Q(x-\mu )+\frac{1}{2}\log \det (Q)+\text{const}\\ & =-\frac{1}{2}{x}^T(-f^{\prime \prime }(\hat{x})+Q)x+x^T(Q\mu +f^{\prime }(\hat{x})-\hat{x}{f}^{\prime \prime }(\hat{x}))+\frac{1}{2}\log \det (Q)+\text{const}\end{array}$$

In the last line we group the quadratic and linear terms.

def _conditional_gaussian_approximation(
      self, x, y, unconditional_gaussian
  ) -> gmrf.Gaussian:
    """The gaussian approximation of P(x | y, params) at the given x."""
    fpp = self.fpp(x, y)
    precision = unconditional_gaussian.precision.add_diag(-fpp)
    linear_part = (
	    # self.precision @ self.mean
        unconditional_gaussian.information_vector + 
        self.fp(x, y) -
        x * fpp
    )
    # Gaussian with mean -fpp + precision
    return gmrf.Gaussian(precision, precision.solve(linear_part))

The constant is then be determined by using the fact that this must be a probability distribution in $x$ (i.e. use the value of the constant from the log probability function of the gaussian with the same quadratic and linear terms).

Evaluate at ${\hat{x}}_0=\mu$ to maximise. Find this $\hat{x}$ with Newton’s method, which amounts to iteratively solving for the maximum in the very same quadratic approximation. Iteratively solve for ${\hat{x}}_{i+1}$ in

$$(Q-f^{\prime \prime }({\hat{x}}_i)){\hat{x}}_{i+1}=Q\mu +f^{\prime }({\hat{x}}_i)-{\hat{x}}_i{f}^{\prime \prime }({\hat{x}}_i)$$

until stabilisation.

def newton_step(x, y, gaussian):
	return self._conditional_gaussian_approximation(x, y, gaussian).mean
 
fpi = jaxopt.FixedPointIteration(newton_step)
x = fpi.run(gaussian.mean, y, gaussian).params

Where can sparsity help?

Looking at the precision matrix again

$$Q(\theta )=\left(\begin{array}{cc}{\Sigma }_b^{-1}& 0\\ 0& R^{-1}\end{array}\right)$$

• Often we make a Markov assumption on ${\Sigma }_b^{-1}$, for example in ICAR model in spatial settings where we only assign covariance to nearest neighbours. This makes the precision matrix sparse.
  • This is the Gaussian Markov Random Field (GMRF)
• Usually the random effect $b$ size dominates the covariates $\beta$.

This means normally $Q$ is sparse. The four steps that are needed in sparse implementation are:

• Adding to the diagonal. $Q-f^{\prime \prime }({\hat{x}}_i)$
• Multiplying a vector $Q\mu$
• Solving linear systems (Newton step)
• Computing $\log \det (Q)$

More details in Dan’s blog and the README of the jax implementation.

Resources used:

• Adam Howes’ thesis chapter on (R-)INLA
• Dan Simpson’s blog
• Junpeng Lao’s attempt in pymc3
• That Stan team paper
• INLA from scratch
• jax implementation

And the relevant pymc issue.