[ANN] Latte.jl: Probabilistic programming for latent Gaussian models

Hi all! Last week I published Latte.jl, a probabilistic programming framework for latent Gaussian models.

What’s a latent Gaussian model?

Many interesting problems in Bayesian inference can be approached via latent Gaussian models. These are hierarchical models with a certain structure: You have some hyperparameters \theta, a latent field x that depends on \theta, and data y that depends on both. Formally, p(y, x, \theta) = p(\theta) p(x \mid \theta) p(y \mid x, \theta). We further assume that p(x \mid \theta) is either Gaussian directly or is amenable to a Laplace approximation through repeated linearisation. Examples of where such models are relevant include spatial disease maps, time series, hierarchical GLMMs, species distribution modelling, Gaussian processes more broadly, etc…

Why Latte?

Latent Gaussian models are highly structured and MCMC often takes a long time to converge on them. Hence, people have developed specialised inference algorithms for these models that leverage this structure. Probably the most famous one is Integrated Nested Laplace Approximation (INLA), which uses Laplace approximations on the latent field + quadrature on the hyperparameters to get accurate posteriors deterministically and at low runtime. Then there’s TMB, which has a much more flexible (template-based) modelling language, and uses the delta method to get uncertainties. Another method combines Laplace over the latent with HMC over the hyperparameters.

These methods share a lot of machinery, yet at the moment they live in fractured packages (mostly in R), fractured communities, and even fractured lingo.
My goal with Latte is to provide a unified PPL framework (based on DynamicPPL.jl <3) that gives you simplicity and flexibility in your modelling, and makes it easy for you to switch between these different inference methods.
Currently, Latte offers INLA (with many of its tricks), the TMB-style delta method, and HMC-Laplace. You can also convert a Latte model to a DPPL model to hand off to Turing.jl.

Usage example

using Latte, Distributions, LinearAlgebra

@latte function hier_poisson(y, x, group, n_groups)
    τ_u ~ Gamma(2.0, 1.0)
    β   ~ MvNormal(zeros(2), 100.0 * I(2))
    u   ~ IIDModel(n_groups, constraint = :sumtozero)(τ = τ_u)
    for i in eachindex(y)
        y[i] ~ Poisson(exp(β[1] + β[2] * x[i] + u[group[i]]))
    end
end

lgm = hier_poisson(y, x, group, n_groups)

r_inla = inla(lgm, y)          # grid integration over θ (default)
r_tmb  = tmb(lgm, y)           # MAP + Laplace (delta-method) covariance
r_hmc  = hmc_laplace(lgm, y)   # NUTS on the Laplace marginal L(θ)

The website has a bunch of tutorials covering a variety of topics, as well as benchmarks comparing to R-INLA.

Outlook

Latte.jl is in v0.1.0. This doesn’t mean it’s unusable, quite the opposite; I’ve had success in throwing it at a bunch of different problems. But of course that doesn’t guarantee that it’ll always work the way it should. So if things break or are suspiciously slow for your use case, or if you have any other kind of feedback, I’d greatly appreciate it.

If you happen to be at ICML next week, I’ll showcase Latte there on Monday in the Probabilistic Numerics tutorial to identify an oil leak in a harbor. I’ll also give a talk at JuliaCon in August that uses Latte for spatial modelling.

That is an amazing initiative and package @timweiland ! :tada: :clap:

I’ve always wanted to try hierarchical INLA models on unstructured meshes. We currently provide support for Gaussian simulation over manifolds in GeoStats.jl, we call it the LindgrenProcess:

For example, here is the code to simulate the latent Gaussian field on a sphere:

using GeoStats
import GLMakie as Mke

# domain of interest
mesh = simplexify(Sphere((0, 0, 0), 1))

# Lindgren process
proc = LindgrenProcess(0.1)

# unconditional simulation
real = rand(proc, mesh, 2)

fig = Mke.Figure(size = (800, 400))
viz(fig[1,1], real[1].geometry, color = real[1].field)
viz(fig[1,2], real[2].geometry, color = real[2].field)
fig

Is there a way to hook into Latte.jl to add other non-Gaussian variables on top? In other words, it would be awesome if we could take a hierarchical model from Latte.jl and use it with our latent LindgrenProcess to simulate more complex structures.

Thanks for the nice feedback @juliohm ! :folded_hands:

If your question is “can Latte do spatial SPDE-based GPs + nonlinear likelihoods”, the answer is yes. Here’s a tutorial that shows how to fit a log-Gaussian Cox process with Latte: Earthquake Intensity: Spatial Modelling with the SPDE Approach | Latte.jl
This uses GaussianMarkovRandomFields.jl’s MaternModel, which implements the same SPDE discretisation by Lindgren et al.. MaternModel however does not support modelling on manifolds yet I think (should be a small code delta, just haven’t gotten around to fiddling with it yet).

If you want to use LindgrenProcess from GeoStats.jl directly in Latte, that’s definitely possible. You basically just wrap it in a LatentModel and tell it to re-use your precision matrix. I told Claude to try this, and it works; see this gist: Latte.jl + GeoStats.jl: A log-Gaussian Cox process on the sphere · GitHub
This draws Poisson observations from a “true” process on the sphere with some offset. INLA then uses those observations to infer the range, offset and the latent field.

Is this related to what you had in mind?

Thank you! I was asking about the second option with general meshes. I had seen the tutorial on spatial processes over grids.

Will digest the gist when I am back to my computer :slight_smile:

Congrats @timweiland, this looks like a really great package! Have just read through the docs and can’t wait to try it out.

Thanks @timweiland, this is a fantastic effort and the docs and especially all the tutorials are a great resource! I tried out a BYM2-Style model for patent applications in Germany and it worked great:

That’s really cool @Jakob, thanks for sharing!! If you ever hit any road bumps while using it, just let me know :slight_smile:

I think the project website is down?

Seems to be up for me, can you try again?

I’m super excited to try this out sometime soon! And what a pretty docs page, I’m in love with the overall design. So cool!

What do you lose by doing this approximation? Compared to plain Turing.jl / MCMC machinery?

The docs section on INLA might be interesting: INLA | Latte.jl

Like you said, it’s an approximation. So at least in theory, if you let Turing / MCMC run for a veeery long time, it’s always going to be more accurate than INLA. The point is of course that you might not want to wait that long, especially for large-scale models. The promise of INLA is that for the model classes it’s designed for, it achieves very accurate posteriors at high speed (often orders of magnitudes faster than MCMC). The original INLA paper has comparisons to MCMC, and there’s other more recent ones. The Validation site of Latte evaluates INLA via SBC on some toy problems. Speed benchmarks are here.

Of course, I’m not an INLA salesman, and there’s no free lunch. INLA is not applicable to any kind of model, only LGMs. And even within the class of LGMs, things can go wrong; for example if your hyperparameter posterior has a nasty shape, or more broadly speaking if you just have many hyperparameters (then quadrature becomes slow).

Which I guess is part of the reason why there have been several variations of the idea of “doing Laplace many times”, all with their own tradeoffs; like TMB, HMC-Laplace, INLA within MCMC, …

For information, I used INLA as a case for learning Claude and got similar performances here: GitHub - haavardhvarnes/INLA.jl · GitHub . This is not my area of expertice so this will not be maintained by me (though have mailed Rue to take this further but dont think they are into Julia at all). Though there might be some ideas in repo which might be useful for people in this thread.

Got similar benchmark results here: Quality vs R-INLA · Julia INLA Ecosystem

Cool work @Haavard_Hvarnes! I’m planning to keep developing Latte, so if you came across any useful insights while building INLA.jl, I’d be happy to chat :slight_smile:

Thx for starting Latte.jl which I am sure will be a success!

Since it is a couple of months ago I dont remember much except being specific of using ecosystem approach with separate packages and leveraging existing Julia packages. Though I made Claude document and store all plans here : INLA.jl/plans at main · haavardhvarnes/INLA.jl · GitHub, which should contain all of the learning going from step to step. Trust these more than my memory :wink:

Ok, I’ve had a chance to play around with Latte a bit and had a couple of questions:

It seems like there’s a requirement that the data be passed in as the first argument to the @latte function, and that they need to be a single 1D vector. Is there a particular reason for this design decision? There are some circumstances (e.g., your age-structured stock assessment model example) where it would be more convenient to have the data stored in a matrix, or some other structure…

Also I was curious about what AD backend(s) are being used under the hood? A quick poke into the source code makes it look like ForwardDiff is being used for the hyperparameters, I assume based on the usual INLA assumption that there will not be very many of them. What’s being used for the optimizing the latent field?

Thanks!

Good questions!

  1. Data doesn’t need to be passed in as the first argument. Data is identified via (is positional argument) & (appears on LHS of a ~ somewhere). Just tested it as a sanity check, and indeed the same SAM model works with Y as the last argument for me.
  2. You can indeed store your data in a matrix and index it that way. (And I should definitely update the SAM tutorial to do that.) The one constraint is that inla wants an AbstractVector, so you’d have to do inla(..., vec(y)) in this case. That’s something I can fix though and I’ll open an issue for it.
  3. For a fixed hyperparameter config, the latent field is optimised via Gauss-Newton. For exponential-family likelihoods this uses closed-form log grads and Hessians; for custom likelihoods it uses sparse AD.
  4. Yes, atm forward-mode AD via ForwardDiff is the “well-supported” backend. This uses custom chain rules for diffing through Gauss-Newton (via the IFT), and through the GMRF logpdf evaluation. In particular, these rules reuse the expensive parts of the computations (the sparse Cholesky and selected inverse), so forward-mode for k hyperparameters here really isn’t quite as expensive as \mathcal{O(k \cdot \text{sparse factorization})}, which might maybe be a fear. Nevertheless reverse-mode should still be the way to go for really large numbers of hyperparameters. GMRFs.jl has the corresponding rules for Mooncake, they’re just not integrated nicely into the “workspace system” that Latte wants. Mostly because I haven’t found the time for it yet; hopefully I can get this done later this summer. I’ll open an issue for this as well in the Latte repo.

Does this help?

This is extremely impressive, thanks for your work and for sharing this @timweiland!

Yes, very helpful, thanks! I guess the docs already make it clear how the observations are identified.

The one constraint is that inla wants an AbstractVector, so you’d have to do inla(..., vec(y)) in this case.

This was more to my question, whether there was any particular reason inla et al. require an AbstractVector. Sounds like this requirement might be able to be relaxed, which would be great.

Again, this is a fantastic effort, can’t wait to see where it goes!