INLA Part 1: Just keep approximating until you get the thing you want
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, , which give rise to some observed values, . I find this convention slightly odd because everywhere else or 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 themselves have parameters and we call that next ‘level’ of parameter a ‘hyperparameter’ and label them . The data generation story is that we first draw some ‘s, which then go into our distributions to give us parameters, , and in turn those go into our likelihood function, .
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
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 and our hyperparameters . With these quantities we can generally compute anything else we care about.
Bayes rule immediately gives us an expression for these quantities
and the marginals are integrals of this joint:
so it doesn’t seem like such a big problem. However the denominator, the normalising constant , is a problem. Evaluating that would require integrating out both the parameters and hyperparameters
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, , are conditionally independent given and that our parameters, , are distributed according to a multivariate Gaussian family:
Where describes the mean and we have parametrised the covariance structure of using a precision matrix, defined as the inverse of the variance-covariance matrix, .
2. Sparsity
INLA requires that our LGMs have a precision matrix which is sparse; a Gaussian 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 , both of which are in the general case but can be made much more efficient if has structure that we can exploit.
What does this mean for our modelling ideas? A zero in our precision matrix, , implies that and are conditionally independent given the other components, . 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 for the covariance, and take without loss of generality (otherwise just replace by throughout).
Sparsity in : If then and are marginally independent.
The characteristic function of is
The joint characteristic function of a pair is obtained by setting every component of to zero except and , which leaves
If the cross term vanishes and the above can be factorised into . A joint characteristic function factorises if and only if the variables are independent, so .
For Gaussians uncorrelated really does mean independent, which isn’t generally true.
Sparsity in : If then and are conditionally independent given all the other components, .
The density is . Expanding the quadratic form and collecting the terms that involve or ,
where collects everything that involves neither nor . Conditioning on means treating those components as constants, so
with containing all of the cross-terms except the ones of interest, and defined similarly. These are both constant once we condition.
The only factor that couples to is the last one, and when this term drops out. The conditional density can be factorised into a function of and a function of only, which is the definition of conditional independence: .
3. Hyperparameter dimensionality
Lastly LGMs require that the dimension of 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
must be a twice differentiable function with a unique maximum and a negative-definite Hessian at that maximum.
Defining
and
we start by Taylor expanding around its maximum
At the mode the first-derivative term is zero and dropping higher-order terms gives us an approximate result
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:
where is the dimensionality of . I’ve labelled an intermediate step as 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 ‘nodes’, , at which we will sample our integrand, then computing weights, , to apply to those values and it is perfect for INLA for the following reason: AGHQ gives us an approximation of the following form
and it has the excellent property that the approximation is exact when is a polynomial of order . 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, : If we have then we can find the mode, , and negative Hessian, , of . We then factor , and make a change of variables to where . In the coordinates the integrand should be close-ish to a standard Gaussian. Nodes in each of the 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 is plenty. This is the first reason why we want a low dimensional because 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, , which we express in terms of things we know, over latent variables:
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…
- Computing the unnormalised hyperparameter posterior using a Laplace approximation
- Computing the latent field marginals : 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
- Integrating out the s by quadrature to get
Hence my tongue-in-cheek title for this post - just keep approximating.
1. Laplace approximation: ‘LA’
Bayes rule for the hyperparameter posterior gives
The denominator is a constant with respect to 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 and this is where we are going to apply our Laplace approximation. We choose to define our function as
and defining and as the mode and negative Hessian of , respectively, the Laplace approximation gives
Both and are functions of so for every value of 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 (i.e. the number of hyperparameters) to a minimum.
Our form for also lets us write down a form for . The first term is constant wrt and drops out. The second term is the log of the Gaussian prior on
the negative Hessian of which is just . 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
Computing the determinant of this (optionally using the sparse structure present) gives us our unnormalised marginal hyperparameter posterior. The factor doesn’t depend on either, so it can be dropped along with and recovered later at the numerical integration step.
2. Nested: ‘N’
Now we will find an expression for . But actual INLA doesn’t give us the joint posterior over , however we can get the marginal posterior for each component: . 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 we remind ourselves that we approximated by a Gaussian allowing us to integrate analytically. With this we can re-write our marginal as:
Nothing fancy has happened here, this is just the Laplace approximation from above, we’ve simply written out explicitly what that approximation says about . Computationally we already have an expression for this from (1).
The marginal of a multivariate Gaussian is found by reading off the mean and variance term so the simplest form for our LGM marginal posterior is given by
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
and we once again apply the Laplace approximation to the integral in the numerator. The numerator is an integral over with both and held fixed, so we can Laplace approximate it. Its mode and negative Hessian now depend on as well as , and the result is
where terms constant wrt 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 and we have a Hessian correction, also a function of , 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 value and every value we have to optimise to find a new mode, and compute a new matrix determinant. The cost is one -dimensional optimisation, and a determinant, for every point on the line, for every node, for every coordinate …it’s a lot. This is going to be pretty slow and hopefully it is now very clear why we demand that 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 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 and standardise so that the Gaussian marginal from (2a) has mean zero and variance one. The log of the full Laplace marginal from (2b), , is some function of which the (2a) Gaussian approximates by . Expanding it to third order gives
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 is Gaussian its third derivative is zero and so only the likelihood contributes:
a vector of scalars we can evaluate once at the mode without an inner optimisation per , and no determinant per . 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 , an unnormalised from step (1) and a conditional marginal from step (2a), (2b) or (2c). However, we actually need the marginal posterior with integrated out
This is an integral over 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 of by optimising, with each evaluation of the objective costing one inner optimisation over , compute the Hessian there, and build the AGHQ nodes with quadrature weights , where , and 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 we dropped in step (1) is recovered by normalising over the nodes,
and the denominator of which is an estimate of itself, the marginal likelihood, which is useful for comparing models. If we dropped the 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,
Phew. That’s it. Laplace over to get , Gaussian or Laplace over to get , quadrature over 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
- Sparse Matrices 1: The linear algebra of linear mixed effects models and their generalisations - Dan Simpson↩
- Fast approximate Bayesian inference - Adam Howes↩
- Bayesian inference with INLA - Virgilio Gómez-Rubio↩
- Bayesian Computing with INLA: A Review - Rue et al.↩
- The Best Of Both Worlds: Hierarchical Linear Regression in PyMC - Thomas Wiecki & Danne Elbers ↩
- Simplified integrated nested Laplace approximation - Simon N. Wood, Biometrika 107(1), 223–230, 2020↩
jax.experimental.sparsemodule↩- pyINLA: the Python interface for INLA↩