All writing

INLA Part 1: Just keep approximating until you get the thing you want

4,090 words Statistics

Deriving the central results required to understand INLA, and getting us ready to implement it for some model building in later posts

Introduction

The ‘Integrated Nested Laplace Approximation’ is a way to fit a subset of hierarchical Bayesian models in a fast and deterministic way.

In the first post of this series we’re going to work through the theory required to understand how we can stitch together approximation after approximation and eventually get some quantities of interest: marginal posterior distributions for our model parameters and hyperparameters!

In the following posts we’ll build an implementation of INLA in python, and in particular using the jax framework which gives us rapid numerical computation and convenient auto-differentiation. In the second post of this series I’ll demonstrate applying INLA on a model where the approximations we require are really good, so good in fact that they give the exact result! Then in the third and probably final post we’ll take a real dataset and this time INLA’s approximations will be put to the test. We’ll take a look at an example of overdispersed count data with the classic Horseshoe Crab dataset and we’ll fit hierarchical Poisson models, without and eventually with, an IID random effect term. I’ll treat the Markov chain Monte Carlo (MCMC) posteriors as the gold standard and compare them to the various INLA approximations we can generate. I’ll also fit the exact same model using pyINLA [8] (a) to check that my implementation is correct (b) to see how much performance I’ve lost by not bothering with all the clever sparse matrix stuff which a proper INLA package has to deal with and (c) to put you off using any of the code I’ll show here for anything other than playing around.

There are lots of good INLA resources out there already, though mostly they are tied to the R-INLA package. I enjoyed reading [1], [2], [3] and [4] while learning about INLA in preparation for writing this.

Hierarchical models

We will be spending our time trying to estimate unobserved model parameters, xx, which give rise to some observed values, yy. I find this convention slightly odd because everywhere else xx or XX would represent features of our model or a design matrix for a model…but conventions are conventions so I’ll stick with it.

Put simply, hierarchical models are Bayesian models where the parameters xx themselves have parameters and we call that next ‘level’ of parameter a ‘hyperparameter’ and label them θ\theta. The data generation story is that we first draw some θ\theta‘s, which then go into our distributions p(x∣θ)p(x|\theta) to give us parameters, xx, and in turn those go into our likelihood function, p(y∣x,θ)p(y | x, \theta).

Diagrammatically this can be expressed as follows:

flowchart LR
  T(["$$\begin{gathered}\theta\\ \text{: hyperparameters}\end{gathered}$$"])
  X(["$$\begin{gathered}x\\ \text{: parameters}\end{gathered}$$"])
  Y(["$$\begin{gathered}y\\ \text{: observations}\end{gathered}$$"])
  T -- "$$p(x \mid \theta)$$" --> X -- "$$p(y \mid x, \theta)$$" --> Y
  T -.-> Y

This is an interesting and powerful class of models. I’ve written a little on hierarchical models in the past so I won’t go in to detail about their advantages here but if you are interested I recommend Thomas Wiecki’s radon post [5] for a friendly introduction.

Mathematically this structure of model can be expressed, using the rules of conditional probability, as follows

p(y,x,θ)=p(y∣x,θ)p(x∣θ)p(θ).p(y, x, \theta) = p(y | x, \theta)p(x| \theta)p(\theta).

Where the terms on the right-hand side are, in order of appearance, our chosen likelihood function, the conditional prior for our model parameters and the prior on the hyperparameters, known as the ‘hyperprior’.

The name of the game with these models is, when given a dataset, to choose an appropriate likelihood, prior and hyperprior and then to find expressions for, or otherwise sample from, the marginal posterior distributions for our parameters p(x∣y)p(x | y) and our hyperparameters p(θ∣y)p(\theta | y). With these quantities we can generally compute anything else we care about.

Bayes rule immediately gives us an expression for these quantities

p(x,θ∣y)=p(y∣x,θ)p(x,θ)p(y)=p(y∣x,θ)p(x∣θ)p(θ)p(y),p(x, \theta | y) = \frac{p(y | x, \theta)p(x, \theta)}{p(y)} = \frac{p(y | x, \theta)p(x | \theta)p(\theta)}{p(y)},

and the marginals are integrals of this joint:

p(θ∣y)=∫p(x,θ∣y) dx,p(xi∣y)=∫∫p(x,θ∣y) dx−i dθ.p(\theta | y) = \int p(x, \theta | y)\,dx, \qquad p(x_i | y) = \int\int p(x, \theta | y)\,dx_{-i}\,d\theta.

so it doesn’t seem like such a big problem. However the denominator, the normalising constant p(y)p(y), is a problem. Evaluating that would require integrating out both the parameters and hyperparameters

p(y)=∫p(y,x,θ)dxdθp(y) = \int p(y, x, \theta)dxd\theta

and in general this cannot be done. MCMC cleverly avoids ever having to compute the normalising constant and is the main way that models like those described here are estimated. MCMC is, however, slow to run, intrinsically random and it can struggle to correctly estimate models with tricky ‘geometries’. When MCMC doesn’t hit those struggles the results are asymptotically exact and so I will be using MCMC posterior distributions as the ‘gold standard’ in this series, where my models differ from the MCMC sampled results, and trust me they will, we will treat those as a result of the approximations INLA requires (or bugs in my implementation!).

Latent Gaussian Models

INLA targets a subset of hierarchical models called Latent Gaussian Models (LGMs). LGMs impose three structural constraints on top of the structure described above.

1. Conditional Independence

We require that our observations, yy, are conditionally independent given (x,θ)(x, \theta) and that our parameters, xx, are distributed according to a multivariate Gaussian family:

p(y∣x,θ)=∏ip(yi∣xi,θ)p(y | x, \theta) = \prod_i p(y_i | x_i, \theta) p(x∣θ)=N(μ(θ),Q−1(θ))p(x | \theta) = \mathcal{N}(\mu(\theta), \mathcal{Q}^{-1}(\theta))

Where μ(θ)\mu(\theta) describes the mean and we have parametrised the covariance structure of xx using a precision matrix, defined as the inverse of the variance-covariance matrix, Σ\Sigma.

2. Sparsity

INLA requires that our LGMs have a precision matrix which is sparse; a Gaussian xx with a sparse precision matrix is called a Gaussian Markov Random Field (GMRF). This is purely for computational efficiency because we are going to end up computing an awful lot of determinants and linear solves involving Q\mathcal{Q}, both of which are O(N3)\mathcal{O}(N^3) in the general case but can be made much more efficient if Q\mathcal{Q} has structure that we can exploit.

What does this mean for our modelling ideas? A zero in our precision matrix, Qij=0\mathcal{Q}_{ij} = 0, implies that xix_i and xjx_j are conditionally independent given the other components, x−(ij)x_{-(ij)}. Practically this means that each component of the latent field only ‘talks to’ a few other components and, luckily, a lot of interesting models fit this description. Random walk ideas, IID random effects models, fixed effects, splines and spatial models like CAR/BYM are all examples of sparse LGMs.

I should point out that I haven’t used sparse matrix ideas in the code shown in these posts. My aim in writing this was to understand and explain how INLA is used to fit the models described above, I wasn’t too fussed about writing the most performant code possible so long as the inference ran in a reasonable amount of time. We will see later that this leads to much slower code than a proper implementation by comparing to pyINLA. Perhaps I’ll revisit this in later posts to see if jax’s experimental ‘sparse’ module can bring performance wins [7]

Optional details: what sparsity in the variance-covariance matrix means, vs. what it means in the precision matrix

Write Σ=Q−1\Sigma = \mathcal{Q}^{-1} for the covariance, and take μ=0\mu = 0 without loss of generality (otherwise just replace xx by x−μx - \mu throughout).

Sparsity in Σ\Sigma: If Σij=0\Sigma_{ij} = 0 then xix_i and xjx_j are marginally independent.

The characteristic function of x∼N(0,Σ)x \sim \mathcal{N}(0, \Sigma) is

φx(t)=E[e i t⊤x]=exp⁡(−12 t⊤Σ t).\varphi_x(t) = \mathbb{E} \left[e^{\,i\, t^\top x}\right] = \exp \left(-\tfrac{1}{2}\, t^\top \Sigma\, t\right).

The joint characteristic function of a pair (xi,xj)(x_i, x_j) is obtained by setting every component of tt to zero except tit_i and tjt_j, which leaves

φxi,xj(ti,tj)=exp⁡(−12[Σiiti2+2 Σij titj+Σjjtj2]).\varphi_{x_i, x_j}(t_i, t_j) = \exp \left(-\tfrac{1}{2}\left[\Sigma_{ii} t_i^2 + 2\,\Sigma_{ij}\, t_i t_j + \Sigma_{jj} t_j^2\right]\right).

If Σij=0\Sigma_{ij} = 0 the cross term vanishes and the above can be factorised into exp⁡(−12Σiiti2)exp⁡(−12Σjjtj2)=φxi(ti) φxj(tj)\exp(-\tfrac12 \Sigma_{ii} t_i^2)\exp(-\tfrac12 \Sigma_{jj} t_j^2) = \varphi_{x_i}(t_i)\,\varphi_{x_j}(t_j). A joint characteristic function factorises if and only if the variables are independent, so xi⊥xjx_i \perp x_j.

For Gaussians uncorrelated really does mean independent, which isn’t generally true.

Sparsity in Q\mathcal{Q}: If Qij=0\mathcal{Q}_{ij} = 0 then xix_i and xjx_j are conditionally independent given all the other components, x−ijx_{-ij}.

The density is p(x)∝exp⁡(−12 x⊤Q x)p(x) \propto \exp \left(-\tfrac{1}{2}\, x^\top \mathcal{Q}\, x\right). Expanding the quadratic form and collecting the terms that involve xix_i or xjx_j,

x⊤Q x=Qii xi2+Qjj xj2+2 Qij xixj+2xi∑k∉{i,j}Qikxk+2xj∑k∉{i,j}Qjkxk+c(x−ij),x^\top \mathcal{Q}\, x = \mathcal{Q}_{ii}\, x_i^2 + \mathcal{Q}_{jj}\, x_j^2 + 2\,\mathcal{Q}_{ij}\, x_i x_j + 2 x_i \sum_{k \notin \{i,j\}} \mathcal{Q}_{ik} x_k + 2 x_j \sum_{k \notin \{i,j\}} \mathcal{Q}_{jk} x_k + c(x_{-ij}),

where c(x−ij)c(x_{-ij}) collects everything that involves neither xix_i nor xjx_j. Conditioning on x−ijx_{-ij} means treating those components as constants, so

p(xi,xj∣x−ij)∝exp⁡(−12[Qiixi2+2xiai])exp⁡(−12[Qjjxj2+2xjaj])exp⁡(−Qij xixj),p(x_i, x_j \mid x_{-ij}) \propto \exp \left(-\tfrac12\left[\mathcal{Q}_{ii} x_i^2 + 2 x_i a_i\right]\right)\exp \left(-\tfrac12\left[\mathcal{Q}_{jj} x_j^2 + 2 x_j a_j\right]\right)\exp \left(-\mathcal{Q}_{ij}\, x_i x_j\right),

with ai=∑k∉{i,j}Qikxka_i = \sum_{k \notin \{i,j\}} \mathcal{Q}_{ik} x_k containing all of the cross-terms except the ones of interest, and aja_j defined similarly. These are both constant once we condition.

The only factor that couples xix_i to xjx_j is the last one, and when Qij=0\mathcal{Q}_{ij} = 0 this term drops out. The conditional density can be factorised into a function of xix_i and a function of xjx_j only, which is the definition of conditional independence: xi⊥xj∣x−ijx_i \perp x_j \mid x_{-ij}.

3. Hyperparameter dimensionality

Lastly LGMs require that the dimension of θ\theta is ‘small’. Now, ‘small’ is ill-defined. [2] recommends “2 to 5, but not exceeding 20”. We’ll see why this is a requirement shortly.

With these three constraints stated we can move on to look at the ideas we’re going to need to actually perform inference on these models.

Laplace approximation

The Laplace approximation is a way to approximately compute integrals of the form

L=∫exp⁡(h(x))dx.\mathcal{L} = \int\exp{(h(x))}dx.

hh must be a twice differentiable function with a unique maximum and a negative-definite Hessian at that maximum.

Defining

x^=arg⁡max⁡xh(x)\widehat{x} = \arg\max_x h(x)

and

H^=−∇x2h(x)∣x=x^,\widehat{\mathcal{H}} = \left. -\nabla^2_{x} h(x)\right|_{x = \widehat{x}},

we start by Taylor expanding hh around its maximum

L=∫exp⁡(h(x^)+(x−x^)⊤∇h(x^)−12(x−x^)⊤H^(x−x^)+…)dx.\mathcal{L} = \int\exp{\Big(h(\widehat{x}) + (x - \widehat{x})^\top \nabla h(\widehat{x}) - \frac{1}{2}(x - \widehat{x})^\top\widehat{\mathcal{H}}(x - \widehat{x}) + \ldots\Big)}dx.

At the mode the first-derivative term is zero and dropping higher-order terms gives us an approximate result

L≈∫exp⁡(h(x^))exp⁡(−12(x−x^)⊤H^(x−x^))dx,=exp⁡(h(x^))∫exp⁡(−12(x−x^)⊤H^(x−x^))dx,\begin{align*} \mathcal{L} \approx& \int\exp{(h(\widehat{x}))} \exp{\Big( - \frac{1}{2}(x - \widehat{x})^\top\widehat{\mathcal{H}}(x - \widehat{x})\Big)}dx,\tag{$\ast$} \\ =& \exp{(h(\widehat{x}))} \int\exp{\Big( - \frac{1}{2}(x - \widehat{x})^\top\widehat{\mathcal{H}}(x - \widehat{x})\Big)}dx, \end{align*} L≈exp⁡(h(x^))(2π)N2∣H^∣12[∣H^∣12(2π)N2∫exp⁡(−12(x−x^)⊤H^(x−x^))dx]\mathcal{L} \approx \exp{(h(\widehat{x}))}\frac{(2\pi)^{\frac{N}{2}}}{|\widehat{\mathcal{H}}|^{\frac{1}{2}}} \Bigg[\frac{|\widehat{\mathcal{H}}|^{\frac{1}{2}}}{(2\pi)^{\frac{N}{2}}} \int\exp{\Big( - \frac{1}{2}(x - \widehat{x})^\top\widehat{\mathcal{H}}(x - \widehat{x})\Big)}dx\Bigg]

where we’ve introduced some factors which cancel out to leave us with the expression in square brackets which we recognise as the integral over a multivariate Gaussian density, which integrates to 1, and leaves us with our result:

L≈exp⁡(h(x^))(2π)N2∣H^∣12\mathcal{L} \approx \exp{(h(\widehat{x}))}\frac{(2\pi)^{\frac{N}{2}}}{|\widehat{\mathcal{H}}|^{\frac{1}{2}}}

where NN is the dimensionality of xx. I’ve labelled an intermediate step as (∗)(\ast) above because I want to draw attention to the fact that this ‘Taylor expand up to 2nd order and drop the linear term’ has resulted in us approximating something as a Gaussian density and we’ll need this idea again later.

Quadrature

The Laplace approximation is great if your integral is in the right form, but how can we approximate the result if it’s not? We turn to numerical integration aka quadrature.

I find quadrature/numerical integration a bit tedious, probably this is a hangover from having to do the trapezium rule with pen and paper at school…But, it’s an important part of INLA so we have to say something about it.

Following [2] we will use Adaptive Gauss Hermite Quadrature (AGHQ). This is a method for selecting the kk ‘nodes’, z1,…,zkz_1, \ldots, z_k, at which we will sample our integrand, then computing weights, w1,…,wkw_1, \ldots, w_k, to apply to those values and it is perfect for INLA for the following reason: AGHQ gives us an approximation of the following form

∫f(z) e−z2/2 dz≈∑j=1kwj f(zj),\int f(z)\,e^{-z^2/2}\,dz \approx \sum_{j=1}^{k} w_j\, f(z_j),

and it has the excellent property that the approximation is exact when f(z)f(z) is a polynomial of order 2k−12k - 1. So if our integrand is fairly close to a standardised multivariate Gaussian then our approximation should be quite good.

The adaptive part of AGHQ simply provides a way to apply this to multivariate Gaussians with a general mean and covariance using an estimation of the curvature at the mode of the integrand, g(θ)g(\theta): If we have dim(θ)=m\textrm{dim}(\theta) = m then we can find the mode, θ^\widehat{\theta}, and negative Hessian, Hθ\mathcal{H}_\theta, of log⁡g(θ)\log g(\theta). We then factor Hθ−1=LL⊤\mathcal{H}_\theta^{-1} = LL^\top, and make a change of variables to zz where θ=θ^+Lz\theta = \widehat{\theta} + Lz. In the zz coordinates the integrand should be close-ish to a standard Gaussian. Nodes in each of the mm dimensions are combined as a product grid.

As far as I can see from the references used here 5 or so nodes in each dimension of θ\theta is plenty. This is the first reason why we want a low dimensional θ\theta because 5dim(θ)5^{\textrm{dim}(\theta)} is fine with 5 hyperparameters (3125) but not great at even 10 hyperparameters (about 9 million!).

INLA

We now have enough to derive expressions for the marginal posteriors of our LMGs. In the hierarchical models section above we saw the two things we want can be expressed as integrals of the joint density, p(y,x,θ)p(y, x, \theta), which we express in terms of things we know, over latent variables:

p(θ∣y)∝∫p(y,x,θ) dx,p(xi∣y)∝∫∫p(y,x,θ) dx−i dθ.p(\theta | y) \propto \int p(y, x, \theta)\,dx, \qquad p(x_i | y) \propto \int\int p(y, x, \theta)\,dx_{-i}\,d\theta.

We have two integrals to compute, and two approaches for computing integrals described so the way forward should be quite clear!

I’ll start with a high level summary of the INLA approach, in 3 steps, and then we’ll work through them. This is important so I’m going to put one of my new fancy blog boxes around it…

INLA in a nutshell
  1. Computing the unnormalised hyperparameter posterior p(θ∣y)p(\theta | y) using a Laplace approximation
  2. Computing the latent field marginals p(xi∣θ,y)p(x_i | \theta, y): There are three ways to do this
    • 2a Gaussian approximation: Least accurate, but fastest
    • 2b Another Laplace Approximation: Most accurate, but slowest
    • 2c The ‘simplified INLA’ approach: Not simpler really! but faster than (2b) while giving similar levels of accuracy
  3. Integrating out the θ\thetas by quadrature to get p(xi∣y)p(x_i | y)

Hence my tongue-in-cheek title for this post - just keep approximating.

1. Laplace approximation: ‘LA’

Bayes rule for the hyperparameter posterior gives

p(θ∣y)=p(y,θ)p(y)=∫p(y,x,θ)dxp(y)∝∫p(y,x,θ)dxp(\theta | y) = \frac{p(y, \theta)}{p(y)} = \frac{\int p(y, x, \theta) dx}{p(y)} \propto \int p(y, x, \theta) dx

The denominator is a constant with respect to θ\theta so we can ignore this term yielding an unnormalised density. In later steps we will use numerical integration to recover a normalised result which we can actually work with.

The numerator is an integral with respect to xx and this is where we are going to apply our Laplace approximation. We choose to define our hh function as

hθ(x)=log⁡p(y,x,θ)=log⁡p(θ)+log⁡p(x∣θ)+∑ilog⁡p(yi∣xi,θ)h_\theta(x) = \log p(y, x, \theta) = \log p(\theta) + \log p(x | \theta) + \sum_i \log p(y_i | x_i, \theta)

and defining x^(θ)\widehat{x}(\theta) and H^(θ)\widehat{\mathcal{H}}(\theta) as the mode and negative Hessian of hθh_\theta, respectively, the Laplace approximation gives

p~(θ∣y)∝exp⁡(hθ(x^(θ)))(2π)N2∣H^(θ)∣12=p(y,x^(θ),θ)(2π)N2∣H^(θ)∣12.\tilde{p}(\theta | y) \propto \exp{\big(h_\theta(\widehat{x}(\theta))\big)} \frac{(2\pi)^{\frac{N}{2}}}{|\widehat{\mathcal{H}}(\theta)|^{\frac{1}{2}}} = p\big(y, \widehat{x}(\theta), \theta\big) \frac{(2\pi)^{\frac{N}{2}}}{|\widehat{\mathcal{H}}(\theta)|^{\frac{1}{2}}}.

Both x^\widehat{x} and H^\widehat{\mathcal{H}} are functions of θ\theta so for every value of θ\theta we have to do a little optimisation problem to find the mode and then compute the Hessian at that mode. This is the second reason that we want to keep the dimensionality of θ\theta (i.e. the number of hyperparameters) to a minimum.

Our form for hh also lets us write down a form for H^(θ)\widehat{\mathcal{H}}(\theta). The first term is constant wrt xx and drops out. The second term is the log of the Gaussian prior on xx

−12(x−μ)⊤Q(θ)(x−μ)-\frac{1}{2}(x - \mu)^\top\mathcal{Q}(\theta)(x - \mu)

the negative Hessian of which is just Q(θ)\mathcal{Q}(\theta). The last term is the Hessian of our log-likelihood, but our conditional independence assumption about the likelihood means that all the off-diagonal terms are zero and we can write

H^(θ)=Q(θ)+diag⁡(c(θ))ci(θ)=−∂2∂xi2log⁡p(yi∣xi,θ)∣xi=x^i(θ).\begin{align*} \widehat{\mathcal{H}}(\theta) &= \mathcal{Q}(\theta) + \operatorname{diag}\big(c(\theta)\big) \\ \quad c_i(\theta) &= -\left.\frac{\partial^2}{\partial x_i^2}\log p(y_i | x_i, \theta)\right|_{x_i = \widehat{x}_i(\theta)}. \end{align*}

Computing the determinant of this (optionally using the sparse structure present) gives us our unnormalised marginal hyperparameter posterior. The (2π)N2(2\pi)^{\frac{N}{2}} factor doesn’t depend on θ\theta either, so it can be dropped along with p(y)p(y) and recovered later at the numerical integration step.

2. Nested: ‘N’

Now we will find an expression for p(x∣θ,y)p(x | \theta, y). But actual INLA doesn’t give us the joint posterior over xx, however we can get the marginal posterior for each component: p(xi∣θ,y)p(x_i | \theta, y). I say ‘actual’ because (2a) will give us a joint posterior, but we haven’t done any nesting so that is more like ‘ILA’.

Here we have three paths we can take each of which trades off approximation accuracy for computational performance.

2a. The Gaussian approximation

The easiest thing we can do is to re-use a result from (1). If we look again at (∗)(\ast) we remind ourselves that we approximated exp⁡(hθ(x))=p(y,x,θ)\exp(h_\theta(x)) = p(y, x, \theta) by a Gaussian allowing us to integrate analytically. With this we can re-write our marginal as:

p(x∣θ,y)=p(y,x,θ)p(y,θ)=p(y,x,θ)∫p(y,x,θ)dx≈exp⁡(h(x^))exp⁡(−12(x−x^)⊤H^(x−x^))∫exp⁡(h(x^))exp⁡(−12(x−x^)⊤H^(x−x^))dx=∣H^∣12(2π)N2exp⁡(−12(x−x^)⊤H^(x−x^))=N(x ; x^(θ), H^(θ)−1).\begin{align*} p(x | \theta, y) &= \frac{p(y, x, \theta)}{p(y, \theta)}\\ &= \frac{p(y, x, \theta)}{\int p(y, x, \theta)dx} \\ &\approx \frac{\exp{(h(\widehat{x}))} \exp{\Big( - \frac{1}{2}(x - \widehat{x})^\top\widehat{\mathcal{H}}(x - \widehat{x})\Big)}}{\int \exp{(h(\widehat{x}))} \exp{\Big( - \frac{1}{2}(x - \widehat{x})^\top\widehat{\mathcal{H}}(x - \widehat{x})\Big)}dx} \\ &= \frac{|\widehat{\mathcal{H}}|^{\frac{1}{2}}}{(2\pi)^{\frac{N}{2}}} \exp{\Big( - \frac{1}{2}(x - \widehat{x})^\top\widehat{\mathcal{H}}(x - \widehat{x})\Big)} \\ &= \mathcal{N}\big(x\,;\ \widehat{x}(\theta),\ \widehat{\mathcal{H}}(\theta)^{-1}\big). \end{align*}

Nothing fancy has happened here, this is just the Laplace approximation from above, we’ve simply written out explicitly what that approximation says about p(x∣θ,y)p(x | \theta, y). Computationally we already have an expression for this from (1).

The marginal of a multivariate Gaussian is found by reading off the ithi^{\mathrm{th}} mean and variance term so the simplest form for our LGM marginal posterior is given by

p~(xi∣θ,y)=N(xi ; x^i(θ), [H^(θ)−1]ii).\tilde p(x_i | \theta, y) = \mathcal{N}\Big(x_i\,;\ \widehat{x}_i(\theta),\ \big[\widehat{\mathcal{H}}(\theta)^{-1}\big]_{ii}\Big).

We haven’t bothered with the ‘nested’ part of INLA, we’ve just used the same Gaussian approximation and, as we will see in later posts, the accuracy of the resulting posterior will suffer.

2b. The full Laplace approximation

In the ‘full’ Laplace approximation we start with Bayes rule for our desired result

p(xi∣θ,y)=∫p(y,x,θ)dx−ip(y,θ)p(x_i | \theta, y) = \frac{\int p(y, x, \theta) dx_{-i}}{p(y, \theta)}

and we once again apply the Laplace approximation to the integral in the numerator. The numerator is an integral over x−ix_{-i} with both θ\theta and xix_i held fixed, so we can Laplace approximate it. Its mode x^−i(xi,θ)\widehat{x}_{-i}(x_i, \theta) and negative Hessian H^−i(xi,θ)\widehat{\mathcal{H}}_{-i}(x_i, \theta) now depend on xix_i as well as θ\theta, and the result is

p~LA(xi∣θ,y)∝p(y,xi,x^−i(xi,θ),θ) ∣H^−i(xi,θ)∣−12,\tilde{p}_{\text{LA}}(x_i | \theta, y) \propto p\big(y, x_i, \widehat{x}_{-i}(x_i, \theta), \theta\big)\,\big|\widehat{\mathcal{H}}_{-i}(x_i, \theta)\big|^{-\frac{1}{2}},

where terms constant wrt xix_i have been dropped since we will need to compute the normalisation constant in a moment anyway.

Unlike in the (2a) case the mode of the distribution is now a function of xix_i and we have a Hessian correction, also a function of xix_i, and we will hopefully see later that these features lead to accurate posteriors when compared to the MCMC results.

The downside is that for each xix_i value and every θ\theta value we have to optimise to find a new mode, x^−i(xi,θ)\widehat{x}_{-i}(x_i, \theta) and compute a new matrix determinant. The cost is one (N−1)(N-1)-dimensional optimisation, and a determinant, for every point on the xix_i line, for every θ\theta node, for every coordinate ii…it’s a lot. This is going to be pretty slow and hopefully it is now very clear why we demand that θ\theta is low-dimensional and that matrices are as sparse as possible to make computing determinants and solves as efficient as possible.

2c. Simplified Laplace

The full Laplace approximation is accurate because for each value of xix_i it re-estimates the mode of the other coordinates and estimates the local curvature there. These evaluations are also why the full Laplace approximation is the slowest option. The ‘simplified Laplace’ approximation [4] (see also: ‘simplified INLA’ [6]) tries to have its cake and eat it: we never do the inner optimisation and instead Taylor expand around the Gaussian expression we derived in (2a) to try and correct both the location and shape of the posterior distribution.

We fix θ\theta and standardise xix_i so that the Gaussian marginal from (2a) has mean zero and variance one. The log of the full Laplace marginal from (2b), log⁡p~LA(xi∣θ,y)\log \tilde{p}_{\text{LA}}(x_i | \theta, y), is some function of xix_i which the (2a) Gaussian approximates by −12xi2-\frac{1}{2}x_i^2. Expanding it to third order gives

log⁡p~SLA(xi∣θ,y)=const+γi(1)xi−12xi2+16γi(3)xi3+….\log \tilde{p}_{\text{SLA}}(x_i | \theta, y) = \text{const} + \gamma^{(1)}_i x_i - \tfrac{1}{2}x_i^2 + \tfrac{1}{6}\gamma^{(3)}_i x_i^3 + \ldots.

The linear term acts to shift the mode of the distribution and the cubic term allows an asymmetry (i.e. skew) which will adjust the shape. Because the prior on xx is Gaussian its third derivative is zero and so only the likelihood contributes:

∂3hθ∂xj3=∂3∂xj3log⁡p(yj∣xj,θ),\frac{\partial^3 h_\theta}{\partial x_j^3} = \frac{\partial^3}{\partial x_j^3}\log p(y_j | x_j, \theta),

a vector of scalars we can evaluate once at the mode without an inner optimisation per xix_i, and no determinant per xix_i. The cubic on its own isn’t a density, so [4] match the expansion to a skew-normal distribution. [6] gives a different derivation of a similar third-order correction which is easier to implement, and that is the one I’ll use in later posts.

3. Integrated: ‘I’

We now have, for any fixed θ\theta, an unnormalised p~(θ∣y)\tilde{p}(\theta | y) from step (1) and a conditional marginal p~(xi∣θ,y)\tilde{p}(x_i | \theta, y) from step (2a), (2b) or (2c). However, we actually need the marginal posterior with θ\theta integrated out

p(xi∣y)=∫p(xi∣θ,y) p(θ∣y) dθ.p(x_i | y) = \int p(x_i | \theta, y)\, p(\theta | y)\, d\theta.

This is an integral over θ\theta and we haven’t assumed any specific Gaussian form for our hyperparameters so we will have to use quadrature to proceed. We find the mode θ^\widehat{\theta} of log⁡p~(θ∣y)\log\tilde{p}(\theta | y) by optimising, with each evaluation of the objective costing one inner optimisation over xx, compute the Hessian there, and build the K=kmK = k^m AGHQ nodes θ1,…,θK\theta_1, \ldots, \theta_K with quadrature weights Δk=∣L∣ wk ezk⊤zk/2\Delta_k = |L|\, w_k\, e^{z_k^\top z_k / 2}, where LL, zkz_k and wkw_k are familiar from the quadrature section and the exponential undoes the Gaussian weight built into the AGHQ expression.

The normalisation constants we have been dropping as we’ve gone are recovered through this integration. The constant p(y)p(y) we dropped in step (1) is recovered by normalising over the nodes,

wk=p~(θk∣y) Δk∑jp~(θj∣y) Δj,w_k = \frac{\tilde{p}(\theta_k | y)\,\Delta_k}{\sum_{j} \tilde{p}(\theta_j | y)\,\Delta_j},

and the denominator of which is an estimate of p(y)p(y) itself, the marginal likelihood, which is useful for comparing models. If we dropped the (2π)N/2(2\pi)^{N/2} factor as in the text above we shouldn’t forget to account for that (guess who forgot to account for that…).

The latent marginals are then weighted sums over the nodes,

p~(xi∣y)=∑k=1Kwk p~(xi∣θk,y).\tilde{p}(x_i | y) = \sum_{k=1}^{K} w_k\, \tilde{p}(x_i | \theta_k, y).

Phew. That’s it. Laplace over xx to get p(θ∣y)p(\theta | y), Gaussian or Laplace over x−ix_{-i} to get p(xi∣θ,y)p(x_i | \theta, y), quadrature over θ\theta to put them together.

Summary

Part 1 done! We’ve had a brief look at LGMs and specifically LGMs where our field is a GMRF. Armed with some approximate integration techniques (the Laplace approximation and numerical quadrature) we were able to get expressions for the marginal posterior of our latent field as well as our hyperparameters.

In the next post we’ll build an implementation of the ideas we’ve seen here. It’ll be a relatively easy model with some made up data, and the INLA approximations will work very well. Even still we should be able to start to see the advantages and disadvantages of our three approaches to computing the marginal latent field posteriors.

References

  1. Sparse Matrices 1: The linear algebra of linear mixed effects models and their generalisations - Dan Simpson↩
  2. Fast approximate Bayesian inference - Adam Howes↩
  3. Bayesian inference with INLA - Virgilio Gómez-Rubio↩
  4. Bayesian Computing with INLA: A Review - Rue et al.↩
  5. The Best Of Both Worlds: Hierarchical Linear Regression in PyMC - Thomas Wiecki & Danne Elbers ↩
  6. Simplified integrated nested Laplace approximation - Simon N. Wood, Biometrika 107(1), 223–230, 2020↩
  7. jax.experimental.sparse module↩
  8. pyINLA: the Python interface for INLA↩
Search posts