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
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θ+dz 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
logLALCS(θ)=logp(D,z^∣θ)+2dzlog(2π)−ℓH
Algorithm 1 ALCS: marginalised log-likelihood evaluation at θ
Figure 1: Supernova cosmology scaling for Λ CDM and w CDM. (a–b) Test 1: wall time and evidence error δlogZ vs N at fixed dz,block=2 . (c–d) Test 2: wall time (log–log) and δlogZ vs total latent dimension D at fixed N=100 ; the D3 full-NS reference is projected. Top axis in (c) shows latents per object.
Figure 2: Student- t extension ( νtrue=5 , Nobj∈{10,20,50,100,150} ). (a) Evidence error δlogZ 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(zj∣θ,xj) for θ∈{−1,0,+1,+2} . At θ<0 the posterior is near-Gaussian; as θ increases tanh(zj) saturates and flat shoulders develop.
Figure 4: Tanh funnel ( J=10 ). (a) Pointwise evidence error vs θ : near zero for θ<0 , growing to ∼20 nats by θ=3 . (b) IS ESS /K at each θ ( K=5000 ): ≈1 for θ<0 , dropping to ≪0.01 for θ>0 , localising the approximation failure.
Model
dθ
dz
δ
σ^
ESS /K
Ref
Eight Schools
2
8
+0.08±0.04
0.05
1.00
analytic
Radon ( J=85 )
4
85
+0.00±0.16
0.17
1.00
analytic
Brownian Motion ( T=50 )
1
50
+0.06±0.07
0.07
1.00
Kalman
LGCP ( M=100 )
2
100
+0.12±0.07
0.08
0.71
full NS
SV, SP500 ( T=100 )
3
100
−0.32±0.18
0.14
0.67
full NS
IRT ( Ns=400 ) †
1
500
—
—
0.10
Stan †
Table 1: ALCS on six inference gym benchmarks. δ=logZALCS−logZref (mean ± 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; θ posterior validated against Stan MCMC instead.
Model
dz
Param.
Min-ESS
tNUTS (s)
tALCS (s)
Speedup
Eight Schools
8
NC
1149
44
3
15×
Radon ( J=85 )
85
NC
3367
62
4
17×
Brownian ( T=50 )
50
centred
2280
52
4
13×
LGCP ( M=100 )
100
centred
7994
91
33
03×
SV ( T=100 )
100
NC
156
1160
329
04×
IRT ( Ns=400 )
500
centred
314
102
3480
0.03׆
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 logZ .
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
T
dz
logZALCS
σ^
tstep (s)
ttotal (s)
100
100
−393.80
0.14
43
333±13
500
500
−1996.59
0.17
56
644±26
1000
1000
−3926.71
0.16
66
936±11
2516
2516
−10510.17
0.18
85
1831±36
Appendix
Table 3: SV tridiagonal ALCS scaling on SP500 data ( nlive=300 ). tstep : wall time per NS step; ttotal : mean ± std over 4 seeds. The ∼2× increase in tstep for a 25× increase in T reflects GPU parallelism hiding the O(T) tridiagonal cost.
Figure 5: IRT μability posterior: ALCS (blue) vs full joint NUTS (orange, 501D, 4 chains). The marginalised posterior is unbiased despite ESS /K=0.10 .
Nobj
logZ^NS
σ^
δALCS
δStudent
ESS /KALCS
ESS /KStudent
10
−22.03±0.21
0.11
−0.19
−0.04
0.86
0.90
20
−44.70±0.18
0.12
−0.09
+0.08
0.86
0.88
50
−98.79±0.13
0.10
−0.97
−0.10
0.38
0.49
100
−195.21±0.22
0.11
−1.58
+0.12
0.16
0.25
150
−293.94±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. logZ^NS : mean ± cross-seed std over 4 NS initialisations (data fixed at seed 42). σ^ = mean within-run NS bootstrap uncertainty on logZ^NS . δ = error relative to analytic reference (nats). ESS /K is the median across M=200 posterior θ samples with K=5000 IS draws per sample.
n
tfull (ms)
ttri (ms)
Speedup
100
3.57
3.34
1.1 ×
500
5.49
3.52
1.6 ×
1000
12.67
3.59
3.5 ×
2000
45.01
3.67
12 ×
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
tfull (ms)
ttri (ms)
Speedup
Mem. full
Mem. tri
32
0.36
0.72
0.5 ×
2.6 MB
51 KB
512
1.54
0.75
2.1 ×
41 MB
815 KB
2048
6.68
0.79
8.4 ×
164 MB
3.3 MB
8192
26.3
1.03
26 ×
655 MB
13 MB
32768
104
1.92
54 ×
2.6 GB
52 MB
Appendix
Table 6: Vmapped Hessian + log-determinant pipeline ( dz=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 ≈3.5ms (3 forward passes) for all dz tested; the full jax.hessian (red) grows as dz×TL and runs out of memory at dz≥5000 due to the O(dz2) compilation graph for the 200-step scan. The tridiagonal method continues to dz=20,000 (4.4 ms) with no memory pressure. (b) Wall time for the vmapped log-determinant step as a function of batch size ( dz=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×100 matrices) but saturates as the GPU parallelises across the batch dimension. The crossover occurs around batch ≈512 ; by batch =32,768 the tridiagonal logdet is 54× faster and uses 51× less memory (52 MB vs 2.6 GB).
Figure 7: Decomposition of ALCS wall time for Test 1 (supernova cosmology, Λ CDM, scaling N ). (a) Number of dead points Ndead 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θ=2 dimensions is identical; Ndead grows as alogN+b (dashed fit), consistent with the theoretical DKL(θ)∼2dθlogN prediction. (b) Per-likelihood-call cost tsingle vs. N . ALCS remains flat at ≈1.2ms 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 ≈0.15ms (lower because no Hessian computation is needed). The growth in total wall time with N is therefore entirely attributable to Ndead growing with DKL , not to any increase in the per-call computational cost.