Automatic Laplace Collapsed Sampling: Scalable Marginalisation of Latent Parameters via Automatic Differentiation

Toby Lovick; David Yallup; Will Handley

cs.LGMar 27, 2026

Automatic Laplace Collapsed Sampling: Scalable Marginalisation of Latent Parameters via Automatic Differentiation

Authors: Toby Lovick, David Yallup, Will Handley

Organizations: Institute of Astronomy, University of Cambridge · Kavli Institute for Cosmology, University of Cambridge · Institute of Astronomy & Kavli Institute for Cosmology, University of Cambridge

Abstract

We present Automatic Laplace Collapsed Sampling (ALCS), a general framework for marginalising latent parameters in Bayesian models using automatic differentiation, which we combine with nested sampling to explore the hyperparameter space in a robust and efficient manner. At each nested sampling likelihood evaluation, ALCS collapses the high-dimensional latent variables $z$ to a scalar contribution via maximum a posteriori (MAP) optimisation and a Laplace approximation, both computed using autodiff. This reduces the effective dimension from $d_θ+ d_z$ to just $d_θ$ , making Bayesian evidence computation tractable for high-dimensional settings without hand-derived gradients or Hessians, and with minimal model-specific engineering. The MAP optimisation and Hessian evaluation are parallelised across live points on GPU-hardware, making the method practical at scale. We also show that automatic differentiation enables local approximations beyond Laplace to parametric families such as the Student- $t$ , which improves evidence estimates for heavy-tailed latents. We validate ALCS on a suite of benchmarks spanning hierarchical, time-series, and discrete-likelihood models and establish where the Gaussian approximation holds. This enables a post-hoc ESS diagnostic that localises failures across hyperparameter space without expensive joint sampling.

Figures & tables

\log\mathcal{L}_{\mathrm{ALCS}}(\bm{\theta})=\log p(\mathcal{D},\hat{\bm{z}}\mid\bm{\theta})+\tfrac{d_{z}}{2}\log(2\pi)-\ell_{H}

Algorithm 1 ALCS: marginalised log-likelihood evaluation at

\bm{\theta}

Figure 1: Supernova cosmology scaling for

\Lambda

CDM and

w

CDM. (a–b) Test 1: wall time and evidence error

\delta\log\mathcal{Z}

N

at fixed

d_{z,\text{block}}=2

. (c–d) Test 2: wall time (log–log) and

\delta\log\mathcal{Z}

vs total latent dimension

D

at fixed

N=100

; the

D^{3}

full-NS reference is projected. Top axis in (c) shows latents per object.

Figure 2: Student-

t

extension (

\nu_{\text{true}}=5

N_{\text{obj}}\in\{10,20,50,100,150\}

). (a) Evidence error

\delta\log\mathcal{Z}

for Gaussian (grey) and Student-

t

(blue) ALCS. (b) IS ESS

/K

(median over

M{=}200

posterior samples). See Table 4 for numerical values.

Figure 3: Latent conditional

p(z_{j}\mid\theta,x_{j})

for

\theta\in\{-1,0,+1,+2\}

. At

\theta<0

the posterior is near-Gaussian; as

\theta

increases

\tanh(z_{j})

saturates and flat shoulders develop.

Figure 4: Tanh funnel (

J=10

). (a) Pointwise evidence error vs

\theta

: near zero for

\theta<0

, growing to

{\sim}20

nats by

\theta=3

. (b) IS ESS

/K

at each

\theta

(

K{=}5000

{\approx}1

for

\theta<0

, dropping to

{\ll}0.01

for

\theta>0

, localising the approximation failure.

Model	$d_{\theta}$	$d_{z}$	$\delta$	$\hat{\sigma}$	ESS $/K$	Ref
Eight Schools	2	8	$+0.08\pm 0.04$	$0.05$	$1.00$	analytic
Radon ( $J=85$ )	4	85	$\phantom{+}0.00\pm 0.16$	$0.17$	$1.00$	analytic
Brownian Motion ( $T=50$ )	1	50	$+0.06\pm 0.07$	$0.07$	$1.00$	Kalman
LGCP ( $M=100$ )	2	100	$+0.12\pm 0.07$	$0.08$	$0.71$	full NS
SV, SP500 ( $T=100$ )	3	100	$-0.32\pm 0.18$	$0.14$	$0.67$	full NS
IRT ( $N_{s}=400$ ) †	1	500	—	—	$0.10$	Stan †

Table 1: ALCS on six inference gym benchmarks.

\delta=\log\mathcal{Z}_{\text{ALCS}}-\log\mathcal{Z}_{\text{ref}}

(mean

\pm

std over 5 runs); ESS

/K

is the median IS diagnostic ( Appendix G ) over

M{=}200

posterior samples. † Full joint NS over 501 dimensions is intractable;

\bm{\theta}

posterior validated against Stan MCMC instead.

Model	$d_{z}$	Param.	Min-ESS	$t_{\text{NUTS}}$ (s)	$t_{\text{ALCS}}$ (s)	Speedup
Eight Schools	8	NC	$1149$	44	3	$15\times$
Radon ( $J=85$ )	85	NC	$3367$	62	4	$17\times$
Brownian ( $T=50$ )	50	centred	$2280$	52	4	$13\times$
LGCP ( $M=100$ )	100	centred	$7994$	91	33	$\phantom{0}3\times$
SV ( $T=100$ )	100	NC	$156$	1160	329	$\phantom{0}4\times$
IRT ( $N_{s}=400$ )	500	centred	$314$	102	3480	$0.03\times^{\dagger}$

Table 2: NUTS vs ALCS wall time on the six inference gym benchmarks (4 chains, 2000 warmup + 2000 samples). Min-ESS is the minimum bulk ESS across all parameters (ArviZ). NC = non-centred parameterisation ( Papaspiliopoulos et al., 2007 ) . † Tasks are not directly comparable: NUTS samples the full 501D joint; ALCS marginalises 500 latents and additionally computes

\log\mathcal{Z}

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

$T$	$d_{z}$	$\log\mathcal{Z}_{\text{ALCS}}$	$\hat{\sigma}$	$t_{\text{step}}$ (s)	$t_{\text{total}}$ (s)
100	100	$-393.80$	$0.14$	$43$	$333\pm 13$
500	500	$-1996.59$	$0.17$	$56$	$644\pm 26$
1000	1000	$-3926.71$	$0.16$	$66$	$936\pm 11$
2516	2516	$-10510.17$	$0.18$	$85$	$1831\pm 36$

Appendix

Table 3: SV tridiagonal ALCS scaling on SP500 data (

n_{\text{live}}=300

t_{\text{step}}

: wall time per NS step;

t_{\text{total}}

: mean

\pm

std over 4 seeds. The

{\sim}2\times

increase in

t_{\text{step}}

for a

25\times

increase in

T

reflects GPU parallelism hiding the

\mathcal{O}(T)

tridiagonal cost.

Figure 5: IRT

\mu_{\text{ability}}

posterior: ALCS (blue) vs full joint NUTS (orange, 501D, 4 chains). The marginalised posterior is unbiased despite ESS

/K=0.10

$N_{\text{obj}}$	$\log\hat{\mathcal{Z}}_{\text{NS}}$	$\hat{\sigma}$	$\delta_{\text{ALCS}}$	$\delta_{\text{Student}}$	ESS $/K_{\text{ALCS}}$	ESS $/K_{\text{Student}}$
10	$-22.03\pm 0.21$	$0.11$	$-0.19$	$-0.04$	$0.86$	$0.90$
20	$-44.70\pm 0.18$	$0.12$	$-0.09$	$+0.08$	$0.86$	$0.88$
50	$-98.79\pm 0.13$	$0.10$	$-0.97$	$-0.10$	$0.38$	$0.49$
100	$-195.21\pm 0.22$	$0.11$	$-1.58$	$+0.12$	$0.16$	$0.25$
150	$-293.94\pm 0.10$	$0.11$	$-1.80$	$+0.61$	$0.12$	$0.20$

Appendix

Table 4: Student-

t

prior experiment: evidence estimates and IS ESS

/K

diagnostic.

\log\hat{\mathcal{Z}}_{\text{NS}}

: mean

\pm

cross-seed std over 4 NS initialisations (data fixed at seed 42).

\hat{\sigma}

= mean within-run NS bootstrap uncertainty on

\log\hat{\mathcal{Z}}_{\text{NS}}

\delta

= error relative to analytic reference (nats). ESS

/K

is the median across

M{=}200

posterior

\bm{\theta}

samples with

K{=}5000

IS draws per sample.

$n$	$t_{\text{full}}$ (ms)	$t_{\text{tri}}$ (ms)	Speedup
100	3.57	3.34	1.1 $\times$
500	5.49	3.52	1.6 $\times$
1000	12.67	3.59	3.5 $\times$
2000	45.01	3.67	12 $\times$
5000	oom	3.92	—

Appendix

Table 5: Tridiagonal vs. full Hessian wall time (single evaluation, 200-step ODE likelihood, H200). OOM = out of memory.

Batch	$t_{\text{full}}$ (ms)	$t_{\text{tri}}$ (ms)	Speedup	Mem. full	Mem. tri
32	0.36	0.72	0.5 $\times$	2.6 MB	51 KB
512	1.54	0.75	2.1 $\times$	41 MB	815 KB
2048	6.68	0.79	8.4 $\times$	164 MB	3.3 MB
8192	26.3	1.03	26 $\times$	655 MB	13 MB
32768	104	1.92	54 $\times$	2.6 GB	52 MB

Appendix

Table 6: Vmapped Hessian

+

log-determinant pipeline (

d_{z}=100

latents per object, H200). Memory is for Hessian storage only.

Figure 6: Tridiagonal Hessian benchmarks on H200 GPU. (a) Wall time vs.

n

for the expensive likelihood (200-step ODE scan). The tridiagonal method (blue) saturates at

{\approx}\,3.5\,\mathrm{ms}

(3 forward passes) for all

d_{z}

tested; the full jax.hessian (red) grows as

d_{z}\times T_{\mathcal{L}}

and runs out of memory at

d_{z}\geq 5000

due to the

\mathcal{O}(d_{z}^{2})

compilation graph for the 200-step scan. The tridiagonal method continues to

d_{z}=20{,}000

(4.4 ms) with no memory pressure. (b) Wall time for the vmapped log-determinant step as a function of batch size (

d_{z}=100

latents per object). Full slogdet (red) grows roughly linearly with batch; the tridiagonal logdet (blue) is slower at small batch (the jax.lax.scan Cholesky recurrence has higher overhead than cuSOLVER for a handful of

100\times 100

matrices) but saturates as the GPU parallelises across the batch dimension. The crossover occurs around batch

{\approx}\,512

; by batch

=32{,}768

the tridiagonal logdet is

54\times

faster and uses

51\times

less memory (52 MB vs 2.6 GB).

Figure 7: Decomposition of ALCS wall time for Test 1 (supernova cosmology,

\Lambda

CDM, scaling

N

). (a) Number of dead points

N_{\mathrm{dead}}

as a function of

N

on a log–linear scale, for both ALCS and analytic (exact) marginalisation. The two methods require the same number of dead points, confirming that the outer nested sampling in

d_{\theta}=2

dimensions is identical;

N_{\mathrm{dead}}

grows as

a\log N+b

(dashed fit), consistent with the theoretical

D_{\mathrm{KL}}(\bm{\theta})\sim\tfrac{d_{\theta}}{2}\log N

prediction. (b) Per-likelihood-call cost

t_{\mathrm{single}}

vs.

N

. ALCS remains flat at

{\approx}\,1.2\,\mathrm{ms}

across the full range: the GPU is saturated by the jax.vmap over

N

objects from

N=64

onwards, so the Hessian and MAP optimisation cost is independent of

N

. The analytic marginalisation is likewise flat at

{\approx}\,0.15\,\mathrm{ms}

(lower because no Hessian computation is needed). The growth in total wall time with

N

is therefore entirely attributable to

N_{\mathrm{dead}}

growing with

D_{\mathrm{KL}}

, not to any increase in the per-call computational cost.

Bayesian Inference

Bayesian inference updates probability distributions over unknowns with evidence to quantify uncertainty and guide prediction and decisions.

See all 315 papers

May 19, 2026

Corrected Integrated Laplace Approximation for Bayesian Inference in Latent Gaussian Models

Latent Gaussian models (LGMs) are a popular class of Bayesian hierarchical models that include Gaussian processes, as well as certain spatial models and mixed-effect models. Efficient Bayesian inference of LGMs often requires marginalizing out the latent variables. For LGMs with a non-Gaussian likelihood, exact marginalization is not possible and a popular approach is to do approximate marginalization with an integrated Laplace approximation (ILA). Using ILA produces an approximate posterior which, in some settings, can differ significantly from the correct posterior, which impacts downstream applications. We propose an importance sampling scheme to correct the error introduced by ILA. By increasing the number of samples in importance sampling, the posterior with ILA converges to the correct posterior. This idea is realized with various techniques, including pseudo-marginalization, quasi-Monte Carlo and randomized quasi-Monte Carlo. We implement our methods in an automatic differentiation framework to support gradient-based algorithms when doing inference on the hyperparameters. For the latter, we specifically consider the use of Hamiltonian Monte Carlo. We demonstrate the benefits of reduced error in various applied models.

Jinlin Lai, Charles C. Margossian, Daniel R. Sheldon

Manning College of Information and Computer Sciences University of Massachusetts Amherst · Department of Statistics University of British Columbia

Latent Variable Models Importance Sampling Bayesian Inference Hierarchical Bayesian Modeling

Ask PDF

May 29, 2026

Large-scale Uncertainty Quantification for Latent Variable Models Using Subsampling Markov Chain Monte Carlo

Stochastic gradient Langevin dynamics combined with Gibbs updates (SGLD--Gibbs) provides a highly scalable approach to approximate Bayesian inference in latent variable models. However, it remains unclear how to tune the algorithm's hyperparameters in a principled manner to ensure the uncertainty estimates are statistically meaningful. In this work, we address this gap in tuning guidance by developing a statistical scaling limit theory for SGLD--Gibbs. We derive a joint asymptotic limit for the global parameters and latent variables under appropriate space-time rescaling. We show that global parameters converge to a diffusion-type limit, while each latent variable converges to a jump process, reflecting the use of intermittent Gibbs updates. This joint jump-diffusion structure reveals how latent-variable randomness contributes to the stationary distribution of the global parameters. We leverage our results to propose explicit guidance on hyperparameter tuning for SGLD--Gibbs that ensures meaningful uncertainty quantification. Numerical experiments show that SGLD--Gibbs with our tuning guidance leads to better parameter estimates, uncertainty quantification, and predictive performance than stochastic variational inference.

Xiaoyu Wang, Jonathan H. Huggins

Department of Mathematics & Statistics, Boston University · Faculty of Computing & Data Sciences, Boston University

Markov Chain Monte Carlo Stochastic Gradient Langevin Dynamics Latent Variable Models Uncertainty Quantification Langevin Dynamics Bayesian Inference

Ask PDF

Aug 9, 2026

LazyHMC: Hamiltonian Monte Carlo Simulation for Lazy, Infinite Dimensional Probabilistic Programs

Hamiltonian Monte Carlo (HMC) is a successful generic inference method in probabilistic programming, but in its ordinary formulation it needs gradients and finite-dimensional parameter spaces. In Haskell, lazy evaluation lets probabilistic programs express stochastic processes and other non-parametric Bayesian models over implicit infinite-dimensional spaces. This paper develops new formulations of gradient-based HMC for this infinite-dimensional setting, via lazy evaluation. For automatic differentiation, we provide an analysis based on a new notion of "piecewise analytic under cylindrical analytic partition" (PACAP), to show that even if a program is infinite-dimensional and defined lazily, the gradient of the likelihood function is finitely supported. For the Monte Carlo method itself, we develop several HMC variants and a No-U-Turn Sampler that operate over the infinite-dimensional parameter space but are still productive because of lazy evaluation. Experiments cover Gaussian mixture clustering, random walks, and piecewise-constant regression with Poisson-process changepoints.

Maria-Nicoleta Crăciun, C. -H. Luke Ong, Tom Schrijvers +1

University of Oxford, UK · Nanyang Technological University, Singapore, Singapore · KU Leuven, Leuven, Belgium

Markov Chain Monte Carlo Automatic Differentiation Bayesian Inference Probabilistic Programming

Ask PDF

Abstract

Bayesian Inference

Explore similar work

Corrected Integrated Laplace Approximation for Bayesian Inference in Latent Gaussian Models

Large-scale Uncertainty Quantification for Latent Variable Models Using Subsampling Markov Chain Monte Carlo

LazyHMC: Hamiltonian Monte Carlo Simulation for Lazy, Infinite Dimensional Probabilistic Programs