arXiv CSopen access

ADELIA: Automatic Differentiation for Efficient Laplace Inference Approximations

2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

cloud distributed-computing parallel-computing

distributed computing, parallel computing, cloud

ADELIA: Automatic Differentiation for Efficient Laplace Inference Approximations Afif Boudaoud∗ , Lisa Gaedke-Merzhäuser† , Alexandros Nikolaos Ziogas∗ , Vincent Maillou∗ , Alexandru Calotoiu∗ , Marcin Copik∗ , Håvard Rue† , Mathieu Luisier∗ , Torsten Hoefler∗ ∗ ETH Zurich, Zurich, Switzerland

arXiv:2605.06392v1 [cs.DC] 7 May 2026

Forward: Compact Intermediates

† King Abdullah University of Science & Technology (KAUST), Thuwal, Saudi Arabia

Exploiting Structured Sparsity(a) for Efficient Automatic Differentiation(b) in LargeScale Bayesian Inference: From 2d+1 Evaluations to One Forward-Backward Pass (c) 1 Sparsity-Exploiting AD 1 Up to 7.6× memory reduction 2 Multi-Variate Support 2 Jointly model k correlated variables

3 Multi-GPU Distributed AD Backward: Efficient 3 Concurrent across GPUs via MPI Recomputation (a) Structured Sparsity (b) ADELIA: Our Contributions

Air Pollution Study Speedup↑ Energy↓ Predict pollutant levels from coarse grid data

10 Benchmarks Including 5 productionscale models

4.2×

5.1×

4-7.9×

5–8×

ADELIA: Exact Gradients for Reliable Convergence at a Fraction of the Cost (c) Improvements over SOTA

Fig. 1. ADELIA replaces finite-difference gradients in DALIA [1] with structure-exploiting reverse-mode AD. (a) Exploited sparsity pattern. (b) Key contributions enabling AD at HPC scale. (c) Per-gradient speedup and energy savings on real-world scale models. Abstract—Spatio-temporal Bayesian inference drives environmental and health sciences using latent Gaussian models. Integrated Nested Laplace Approximations (INLA) enable inference for these models at HPC scale but rely on derivative-based optimization over d hyperparameters. State-of-the-art INLA implementations approximate derivatives via central finite differences (FD), requiring 2d+1 evaluations. These evaluations are embarrassingly parallel, but total work and energy grow with d, limiting time-to-solution under fixed budgets. Reverse-mode automatic differentiation (AD) computes exact gradients independently of d, but its efficient application to INLA’s structuredsparse kernels is an open challenge. We present ADELIA, the first AD-enabled INLA implementation with a structure-exploiting multi-GPU backward pass leveraging model sparsity. We evaluate ADELIA on ten benchmark models, including real-world airpollution monitoring. We achieve 4.2–7.9× per-gradient speedups and reliable convergence on production-scale models with up to 1.9M latent variables, where FD struggles. Even when scaled to 16–32 GPUs to match ADELIA’s wall-clock time, FD consumes 5–8× more energy. Index Terms—Automatic Differentiation, INLA

I. I NTRODUCTION Bayesian spatio-temporal models are widely used in climate, environmental, and biomedical sciences [2]–[5] to characterize complex processes and provide principled uncertainty estimates. However, performing inference remains computationally challenging as the spatio-temporal structure typically gives rise to latent parameter space dimensions beyond the feasibility of traditional sampling-based approaches [2], [6]. InteCorresponding author: Afif Boudaoud ([email protected]).

grated Nested Laplace Approximations (INLA) [7] addresses this issue by leveraging the conditional independence structure of the underlying latent Gaussian models, inducing blocktridiagonal arrowhead sparsity (BTA) patterns (Figure 2a) in the associated spatio-temporal precision matrices [8], [9]. INLA’s computational cost is dominated by derivative computations in its optimization phase and during posterior approximation, which requires derivatives with respect to the model’s hyperparameters. State-of-the-art implementations, such as RINLA [7] and DALIA [1], rely on central finite differences (FD), which require 2d+1 function evaluations per gradient for d hyperparameters and 2d2 +1 evaluations for the Hessian during the post-optimization stage, incurring high computational costs and an inherent approximation error. We consider both univariate and multivariate models, where multiple quantities are modeled jointly, raising the number of hyperparameters to d = 15, further exacerbating the computational burden. For scalar-valued functions, reverse-mode automatic differentiation (AD) computes exact gradients at 3–5× the runtime cost of a single forward evaluation, regardless of d [10], [11]. Previous AD-based approaches to the Laplace approximations [12]–[14] use general-purpose sparse solvers on CPU, leaving BTA structure unexploited at HPC scale. For matrices with BTA sparsity, the forward pass preserves this sparsity, but the standard Cholesky derivative is dense, rendering it impractical. Attempting to preserve sparsity with standard AD techniques requires storing intermediate states across loop iterations, exceeding GPU memory at a million-variable scale.

A second challenge is posed by multivariate models, which introduce coupling terms between the k response variables. Here, each block is a weighted combination of k per-variable matrices, and the backward pass must decompose these to compute gradients for both per-variable and cross-variable hyperparameters. Third, the largest models exceed single-GPU capacity, requiring the distribution of both the factorization and backward pass across multiple GPUs. We present ADELIA1 , the first AD-enabled INLA framework for HPC-scale spatio-temporal inference (Figure 1), and evaluate it on ten benchmark models, including real-world applications such as air pollution [1] and temperature [9] modeling. On the four multivariate models we test, FD not only slows convergence but prevents it, stalling at gradient norms orders of magnitude above tolerance; ADELIA’s exact gradients resolve this, achieving well-conditioned optima while providing 4.2–7.9× per-gradient speedups and 5–8× energy savings on equal hardware. These speedups combine algorithmic gains (2d+1 evaluations reduced to one) with framework-level efficiency differences, decomposed in Section IV-G. Our main contributions are: 1) Structure-exploiting gradient computation for INLA. We develop a custom differentiation rule for the full objective function, covering factorization, solve, and selected inversion of BTA precision matrices, that computes exact gradients while exploiting block sparsity. By fusing stages of the forward computation, only compact intermediate matrices need to be stored, reducing memory by approximately 2×. 2) Extension to multivariate models. We extend the gradient computation to jointly model k correlated variables by exploiting the block structure of the multivariate precision matrix, separating per-variable and cross-variable hyperparameter contributions to reduce per-step work (benchmarked up to k=3, d=15). 3) Multi-GPU distributed AD. For models whose stored intermediates exceed single-GPU memory, we distribute the backward pass across multiple GPUs via temporaldomain decomposition with nested dissection reordering. We synchronize boundary intermediates via MPI allgather and stage intermediates on CPU memory, enabling AD on models up to 1.9M latent variables that otherwise cannot run on a single GPU. 4) Comprehensive evaluation on up to 128 GH200 GPUs. Beyond reporting performance results, we provide a practitioner-oriented analysis that disentangles algorithmic gains from framework effects, evaluates when FD can match AD via scaling, and characterizes where time and memory concentrate in the gradient pipeline. II. BACKGROUND A. The INLA Pipeline Environmental monitoring and public health policy rely on spatial and temporal predictions with quantified uncertainty: 1 Code: https://github.com/affifboudaoud/adelia-artifact

an air quality agency must know not only that pollution is elevated, but with what probability it exceeds a regulatory threshold; a climate service must attach confidence intervals to regional projections to inform adaptation planning. Latent Gaussian models [15] provide a principled Bayesian framework for these problems. They consist of two layers of unknowns. The high-dimensional latent field x ∈ RN represents the unobserved spatio-temporal process of interest (e.g., pollution concentration across a fine mesh) and is modeled as a Gaussian Markov random field (GMRF) with sparse precision (inverse covariance) matrix Qp . The second layer of unknowns is a small set of hyperparameters θ ∈ Rd (d ≪ N ) that governs global properties (spatial range, temporal range, observation precision). Integrated nested Laplace approximations (INLA) [7] make inference in these models computationally feasible by replacing the expensive sampling of traditional Bayesian methods (e.g., Markov chain Monte Carlo [2], [6]) with deterministic approximations that exploit the conditional independence structure of the GMRF. This work targets all stages of INLA’s algorithmic pipeline that require derivatives. They comprise two out of the three main stages and account for the majority of the computational workload, namely (1) optimization of the scalar objective function f (θ) over the hyperparameters θ using a quasiNewton method, which requires repeated evaluations and gradients of f ; and (2) computation of the posterior Hessian at the optimum θ ∗ to quantify hyperparameter uncertainty (Section III-D). Both stages rely on the same core linear algebra operations within f , which is defined as: f (θ) := 21 log |Qp | − 12 log |Qc | − 21 x∗T Qp x∗ + log ℓ(y | x∗ , θ) + log p(θ) + const

(1)

Here Qp (θ) is the prior precision matrix and Qc (θ) = Qp (θ) + τ AT A is the precision matrix after conditioning on observations under a Gaussian observation model, where τ is the observation precision, y ∈ Rny denotes the observations, and A ∈ Rny ×N is a sparse projection matrix mapping latent variables to observation locations. The vector x∗ = Q−1 c r, where r = τ AT y, is the posterior mode (the solution of a linear system). The likelihood for the observations is denoted by ℓ(y | x∗ , θ) and p(θ) is the prior for the hyperparameters. Computational Structure: Evaluating f (θ) requires constructing Qp (θ) from sparse Kronecker products of spatial and temporal base matrices [8], [9] and forming Qc ; Cholesky factorizations of both precision matrices for their log-determinants; a triangular solve for x∗ ; and a quadratic form x∗T Qp x∗ . The remaining terms (log ℓ, log p) are cheap to evaluate once x∗ is known. B. BTA structured Sparsity in Spatio-Temporal Models The precision matrices of the employed spatio-temporal models exhibit a block-tridiagonal arrowhead (BTA) sparsity pattern [1], [15], illustrated in Figure 2a. The matrix is composed of n diagonal blocks arising from a temporal discretization over n time steps. Each diagonal block, of size

BT3

BT4

BT5

CT5

CT6

GPU 0

BT2

CT3

CT4

CT2

BT4

BT5

CT5

BT3

CT3

CT4

GPU 1 Shared C1

CT6

(a) BTA matrix partitioned across 2 GPUs (b) Reordering for concurrent execution

D2 D5

GPU 1 GPU 0

BT2

Phase 2

Forward Pass Phase 1'

D2 D5

Phase 2'

GPU 1 GPU 0

Phase 1

Allgather

CT2

CT1

Allreduce

BT1

GPU 1 local GPU 0 local

BT1

Reduced system

CT1

Backward Pass

Fig. 2: Distributed two-phase algorithm for P =2 GPUs on a BTA matrix with n=6 blocks. (a) Original partitioning. (b) Nested dissection reordering into interior chains and reduced system. (c) Forward and backward passes with communication. b×b, is associated with the discretization of the spatial domain at a given time step using b mesh nodes. Di ∈ Rb×b are the main-diagonal blocks, Bi ∈ Rb×b (i = 1, . . . , n−1) are sub-diagonal blocks encoding temporal dependence between consecutive time steps, and Ci ∈ Ra×b together with the arrowhead tip T ∈ Ra×a arise from a fixed-effect (regression) coefficients that couple to every temporal block through the observation model. The conditional precision Qc has full BTA structure, while the prior Qp is block-tridiagonal (BT, i.e., Ci = 0): the arrowhead arises from the observation term τ AT A in Qc . For multivariate models with k correlated response variables, b = k · ns (ns spatial locations), increasing both per-block cost and the number of hyperparameters d = 4k+ k2 (each variable contributes 4 parameters; each pair adds a coupling parameter). In the models we benchmark, n ranges from 5 to 512 temporal steps and b from 92 to 13,455. The arrowhead width ranges from a = 2 to 6, resulting in precision matrix dimensions N = b · n + a. C. Efficient Block Solver For BTA matrices, the Serinv library [16] provides routines for Cholesky factorization in O(n · b3 ) operations, triangular solve, and selected inversion (SI), which computes entries of Q−1 at positions where Q is nonzero, also in O(n · b3 ). Although the total number of latent parameters N ≈ n · b ranges from hundreds to 1.9 million, small by modern machine learning standards, the computational challenge lies in the per-block density: each temporal step requires multiple dense b × b linear algebra routines per optimization step, and the factorization must store n−1 Schur-complement carries of size b×b. For the largest models, the carries alone consume 63 GiB, and together with workspace memory for factorization and triangular solves, require distribution across multiple GPUs (Section III-C). D. Gradient Computation in INLA Both stages of the pipeline described in Section II-A require differentiating f through the BTA factorizations above.

DALIA [1], the state-of-the-art HPC implementation of INLA, has scaled inference to million-parameter models with up to three orders of magnitude speedup over the reference R-INLA implementation [7], [17]. DALIA optimizes f using an LBFGS-B [18] algorithm. This requires the evaluation of f and its gradient in every iteration and computes the gradients via central finite differences: f (θ + hej ) − f (θ − hej ) ∂f , ≈ ∂θj 2h

(2)

where h is a small step size (h=10−3 for gradients, h=5×10−3 for the Hessian) and ej is the j-th unit vector, so each evaluation perturbs a single hyperparameter. The scheme incurs an inherent truncation error of O(h2 ) [18] and requires 2d+1 evaluations of f for d hyperparameters. Automatic differentiation (AD) [10], [11] computes exact derivatives by applying the chain rule to each elementary operation. In reverse mode, the program is evaluated forward while recording a trace of operations; a backward pass then propagates derivatives from the output back through this trace (also called the tape) to accumulate the gradient. Because non-linear operations (e.g., Cholesky factorization) cannot be inverted from their output alone, the backward pass requires intermediate results from the forward pass; these can either be stored during the forward pass and provided to the backward pass, or recomputed when needed, leading to a tradeoff between memory and compute. Unlike forward-mode AD, which scales as O(d), reversemode AD yields the full gradient of a scalar objective in a single backward pass at a proven cost of 3–5× one forward evaluation [10], though realized speedups also depend on perevaluation framework efficiency (Section IV-G). Beyond gradient computation, INLA requires the Hessian of f at the optimum to quantify posterior uncertainty over the hyperparameters. With finite differences, a second-order central difference scheme requires 2d2 +1 function evaluations; Section III-D describes how AD reduces this cost.

Algorithm 1 Fused BTA Cholesky with lower triangular solve. Fused steps are shown in . n Require: BTA blocks {Di }ni=1 , {Bi }n−1 i=1 , {Ci }i=1 of Qc , arrowhead tip T , RHS r n Ensure: Schur carries {Si }n−1 i=1 , lower solve vectors {zi }i=1 , zT , arrow factor LT 1: S0 ← 0, T̃0 ← T , C0L ← 0, LB0 ← 0, z0 ← 0, vT ← 0 2: for i = 1 to n do 3: LDi ← chol(Di − Si−1 ) L LTBi−1 )L−T 4: LCi ← (Ci − Ci−1 Di −1 5: zi ← LDi (ri − LBi−1 zi−1 ), vT ← vT + LCi zi 6: if i < n then 7: LBi ← Bi L−T Si ← LBi LTBi Di , 8: end if 9: T̃i ← T̃i−1 − LCi LTCi , CiL ← LCi 10: end for 11: LT ← chol(T̃n ), zT ← L−1 T (rT − vT )

III. S TRUCTURE -E XPLOITING AD FOR INLA A. Custom Backward Pass for BTA Systems Recall that within INLA’s objective function f (Eq. 1): Qc has BTA structure (Figure 2a), while Qp is block-tridiagonal (BT) without the arrowhead. The dominant cost is computing log |Qc |, log |Qp |, and x∗ via Cholesky factorizations and triangular solves. Naive reverse-mode AD through these operations is impractical: the standard AD rule for differentiating a Cholesky factorization [19] requires L−1 , which is dense even when L is sparse, degrading backward complexity from O(n · b3 ) to O(N 3 ) where N = n · b. Expressing the factorization as a block-by-block loop avoids the dense inverse, but AD must store all intermediate states across iterations for the backward pass, exceeding GPU memory for large models. We address this through three techniques: (a) a fused forward pass that reduces the stored intermediates, (b) carry-based reconstruction that recovers the per-block components of L on the fly during the backward pass from compact stored intermediates, and (c) a hand-derived analytical gradient that decomposes Eq. 1 into independently differentiable terms. 1) Fused Forward Pass: The posterior mode satisfies Qc x∗ = LLT x∗ = r which is solved for x∗ using a lower triangular solve Lz = r and an upper triangular solve LT x∗ = z. Our fused pass (Algorithm 1) combines the Cholesky factorization with the lower triangular solve in one sweep, treating each diagonal factor LDi as a per-step temporary rather than storing the full factor L. Because of this fusion, LDi and LBi are only needed within step i and are discarded after computing the Schurcomplement carry Si (the b×b matrix that propagates coupling from block i to block i+1) and the lower triangular solve vector zi . Only {Si } and {zi } are retained, halving storage relative to the full factor (n vs ∼2n dense b×b blocks).

2) Upper Triangular Solve and L Reconstruction: The second stage LT x∗ = z sweeps i from block n to 1: x∗T = L−T T zT

(3)

T ∗ x∗n = L−T Dn (zn − LCn xT ) T ∗ T ∗ x∗i = L−T Di (zi − LBi xi+1 − LCi xT )

(4) (5)

This, together with the selected inversion (SI; Phase A below) and gradient accumulations (Phases A–C), requires LDi at each block and LBi for i < n. Rather than storing them, we reconstruct on the fly: at block i, the stored carry Si−1 together with the input blocks Di and Bi suffice to recompute LDi via one Cholesky factorization and two triangular solves. This trades compute for memory: one additional Cholesky per block, with peak working memory of O(b2 ). All backward operations share the same sweep direction and reconstructed factors, so they are fused into a single pass over the blocks. 3) Analytical Gradient Decomposition: The objective f (θ) depends on θ both directly (through Qp , Qc , and the priors) and indirectly through the posterior mode x∗ (θ) = Q−1 c r. By ∂f ∂f dx∗ df = + . the chain rule, the total derivative is dθ ∂θk ∂x∗T dθk k The second term vanishes by the envelope theorem [20]: x∗ maximizes f with respect to x, so the stationarity condition ∂f /∂x = 0 holds at x=x∗ , and x∗ can be treated as a constant when differentiating with respect to θ. Differentiating each term of Eq. 1 with x∗ held constant gives three independent gradient contributions, detailed below: ∂f 1 ∂Qc 1 ∂Qp ∗ = − tr Q−1 − x∗T x c ∂θk 2 ∂θk 2 ∂θk {z } | {z } | Phase B Phase A 1 ∂ log p ∂ −1 ∂Qp + tr Qp + + log ℓ(y | x∗ , θ) (6) 2 ∂θk ∂θk ∂θk {z } | Phase C

Each phase is computed analytically from the stored carries {Si } and vectors {zi }, requiring O(n · b3 ) work with O(b2 ) working memory per block. Each factorization produces n carries of size b × b. Since the Qc and Qp sweeps run sequentially, only one set is live at a time, giving a peak memory footprint of n · b2 (carries) plus O(b2 ) per-block workspace for the reconstructed factors, independent of the number of intermediates that naive AD would store. a) Phase A - Selected inversion gradients for log |Qc |: The of a log-determinant is gradient ∂ log |Q| −1 ∂Q = tr Q [15], [21]. Naively, this requires ∂θk ∂θk −1 the full inverse Q , which is dense and costs O(N 3 ) to compute. However, expanding the trace of a product gives P tr(AB) = A Bji , which reduces to an element-wise ij ij sum without forming the full product AB. Since both Q−1 ∂Q and ∂θ are symmetric and the latter has the same BTA k ∂Q sparsity pattern as Q, every term with a zero entry in ∂θ k vanishes, and we only need Q−1 at the nonzero positions of Q: the block-diagonal ZDi = [Q−1 ]ii , block-sub-diagonal ZBi = [Q−1 ]i+1,i , arrowhead column ZCi = [Q−1 ]T,i , and arrowhead tip ZT = [Q−1 ]T,T . These are precisely the

entries computed by Selected Inversion [16] which sweeps backward from block n to 1, recovering ZDi , ZBi , ZCi , and ZT from the Cholesky factors LDi , LBi , LCi , and LT . ∂Bi i The block Jacobians ∂D ∂θk , ∂θk , etc. are inexpensive to obtain from the closed-form block construction. The full trace then reduces to a sum of per-block traces: n n−1 X ∂ log |Qc | X ∂Di T ∂Bi = +2 tr ZDi tr ZBi ∂θk ∂θk ∂θk i=1 i=1 n X ∂T T ∂Ci +2 + tr ZT (7) tr ZCi ∂θk ∂θk i=1 We fuse this trace accumulation into the SI sweep: at each block i, we compute ZDi and ZBi , immediately accumulate their gradient contributions, and discard them, so only two blocks of Q−1 (ZDi and ZDi+1 ) are ever live simultaneously. b) Phase B - Quadratic form gradients: The quadratic form x∗T Qp x∗ expands via the BT block structure as: n n−1 X X ∗ T ∗ x∗T D x + 2 x∗T (8) x∗T Qp x∗ = i i i i Bi xi+1 i=1

i=1

Since x∗ is treated as fixed at the mode, differentiating with respect to θk passes through to the blocks: n n−1 X X ∂Di ∗ ∂Bi ∗ ∂ (x∗T Qp x∗ ) = x∗T x + 2 x∗T x (9) i i i ∂θk ∂θ ∂θk i+1 k i=1 i=1 This requires only the already-computed posterior mode x∗ and the block Jacobians of Qp . c) Phase C - Prior log-determinant gradients: The term log |Qp | has the structure of log |Qc | but uses BT factorization (no arrowhead), so the same carry-based technique applies. While we instantiate this approach for BTA matrices, the underlying pattern generalizes to any block recurrence where (i) a compact carry state summarizes all prior blocks and (ii) per-block factors can be reconstructed from that carry and the local input data. This includes block-tridiagonal, blockbanded, and related elimination-based factorizations. B. Extension to Multivariate Models The previous section developed the backward pass for a single response variable. We now extend it to multivariate spatio-temporal models that jointly model k correlated response variables (e.g., three correlated pollutants in air quality monitoring). The BTA precision matrix retains the same blocktridiagonal pattern, but each diagonal block Dt now groups all k variables at a given time step into a super-block of size b = k · ns , where ns is the number of spatial locations. This super-block is itself a k × k grid of ns × ns sub-blocks, where each sub-block encodes the spatial coupling between a pair of variables. Let Qm (t) ∈ Rns ×ns denote the spatial precision matrix of the m-th variable at time step t, as if that variable were modeled independently. Each (i, j) sub-block of the super-block at time t is a weighted combination of these per-variable matrices: k X (m) (ij) Qsuper (t) = Wij (σ, λ) · Qm (t) (10) m=1

(m)

where Wij are entries of the coregional weight matrix, derived from the linear model of coregionalization [22], parameterized by scale σ and coupling parameters λ, both part of the hyperparameter vector θ, for details see [1] The forward factorization (Section III-A) operates on these super-blocks unchanged. We formulate the super-block as Eq. 10 specifically to exploit the fact that each hyperparameter in θ affects exactly one of the two factors, reducing per-step work in the backward pass. A per-variable hyperparameter θk (e.g., spatial range) affects Qm but not the weights: (ij)

∂Qsuper (m) ∂Qm = Wij · ∂θk ∂θk

(11)

A coupling hyperparameter θk (i.e., one of the σj or λj ) affects the weights but not Qm : (m) k (ij) X ∂Wij ∂Qsuper = · Qm ∂θk ∂θk m=1

(12)

Both enter the SI gradient trace (Eq. 7) through the block i Jacobians ∂D ∂θk . The per-variable case (Eq. 11) requires the (m) m weight Wij and the per-model derivative ∂Q ∂θk ; the coupling case (Eq. 12) requires all k per-model blocks Qm (ti ) themselves. In practice, the Qm are never formed as separate matrices; they share the same underlying spatial components and are combined with the coregionalization weights on the fly during super-block assembly. Since the per-model blocks and Jacobians are reconstructed from the same stored carries and discarded after each step, this adds no persistent memory beyond the existing O(b2 ) per-block workspace. C. Multi-GPU Distributed AD The BTA factorization is inherently block-sequential, where the computation related to each diagonal block depends on the Schur-complement carry from the previous diagonal block. The primary motivation for distribution is memory: for the largest models, the stored carries (n·b2 elements) could exceed the capacity of a single node (GPU and host combined), making distribution a necessity to run these models at all. We adopt Serinv’s two-phase domain decomposition [16] for the forward factorization and extend it to the backward pass, distributing both across P GPUs. On top of enabling larger models, this algorithm also exposes concurrency by reordering the blocks such that the majority of the work (Phase 1) can be done in parallel with no communication, while the remaining work (Phase 2) requires only a single round of communication and redundant local solves. The hand-derived backward pass from Section III-A is essential here: it provides explicit blockby-block control that allows inserting MPI communication and staging carries to CPU memory, which would not be trivially possible within an automatic AD tape. a) Two-Phase Partitioning: The n temporal blocks are divided into P contiguous partitions (Figure 2a). The blocks at partition interfaces (boundary blocks) couple to blocks in both adjacent partitions through the Schur complement carry; the remaining interior blocks couple only within their

own partition. Serinv’s nested dissection reordering [16] (Figure 2b, regions delimited by red lines) separates these two groups: interior blocks stay on their assigned GPU as a selfcontained chain, while boundary blocks are collected into a small shared reduced system of 2P −1 blocks that captures all cross-partition coupling. For example, with n=6 and P =2: GPU 0 keeps D1 , D2 as its interior chain, GPU 1 keeps D5 , and the boundary blocks D3 , D4 , D6 form the reduced system. In Phase 1 (Figure 2c), all GPUs factorize their interior chains concurrently: GPU 0 processes D1 → D2 , GPU 1 processes D5 . In Phase 2, an MPI allgather exchanges boundary data, and each GPU redundantly solves the reduced system (D3 , D4 , D6 ) to avoid communication. For the backward pass, the order reverses: Phase 2′ solves the reduced system first to produce boundary selected inversion values, then Phase 1′ runs concurrent backward sweeps seeded by these values, followed by an allreduce of the d gradient scalars. b) Communication Pattern: Phase 1 requires no interrank communication. After Phase 1, MPI allgathers collect the 2P −1 boundary diagonal (b × b), lower (b × b), and arrow blocks (a × b), plus the lower triangular solve RHS vectors (b × 1), from all ranks; a negligible allreduce sums the arrow tip (a × a) and the interior log-determinant (scalar). In Phase 2, each rank assembles and factorizes the reduced system (2P −1 blocks) locally. A final MPI allreduce of d gradient scalars combines per-rank contributions. The total number of distinct elements exchanged across all collectives is (4P −2)b2 + (2P −1)(a+1)b + a2 + d + 1, dominated by the O(P b2 ) boundary blocks. c) CPU-Staged Schur Carries: The Schur carries (nlocal × b2 elements per rank) are written to CPU memory during the forward pass and read back to GPU one block at a time during the backward pass for L reconstruction. This bounds per-GPU memory to O(b2 ) working memory regardless of nlocal , at the cost of nlocal CPU↔GPU transfers. The same staging is applied to the reduced-system factorization when the 2P −1 blocks cannot be held simultaneously on GPU, keeping the O(b2 ) bound throughout. D. Posterior Hessian via AD Gradients Post optimization, INLA computes the Hessian H = ∇2 f (θ ∗ ) at the optimum to approximate the posterior uncertainty of the hyperparameters. Rather than 2d2 +1 objective function evaluations (second-order FD approach), we compute this Hessian via first-order finite differences of AD gradients: H:,j ≈

∇f (θ ∗ +hej ) − ∇f (θ ∗ −hej ) 2h

(13)

This reduces the cost from 2d2 +1 objective evaluations to 2d + 1 calls to the AD gradient routine. Because the Hessian computation is a one-time post-optimization step, fully secondorder AD would add complexity for a modest gain: the Hessian differentiates the d-vector gradient, not a scalar function, so each of its d columns still requires a separate reverse pass.

E. Implementation Our open-source implementation, ADELIA, extends the DALIA framework [1] with a JAX backend. ADELIA coexists with DALIA’s original FD implementation, which uses CuPy with Serinv’s [16] GPU-optimized BTA solver, allowing direct comparison on identical model configurations. JAX’s jax.custom_vjp primitive registers our structureexploiting backward pass for the BTA factorization, while JAX’s standard reverse-mode AD handles the remaining terms (prior log-probabilities, likelihood); jax.jit compiles each stage via XLA [23] for GPU execution. Multi-GPU communication uses mpi4jax [24] for JAX-compatible MPI operations within XLA-compiled code. IV. E XPERIMENTAL E VALUATION We organize the evaluation around eight questions a practitioner would follow when adopting AD for structured sparse matrix computations. We first establish that structureexploiting AD is necessary, not optional (§IV-B), then validate gradient correctness and show that AD gradients enable convergence where FD stalls (§IV-C). We measure per-gradient speedup over FD on equal hardware (§IV-D), translate these to end-to-end wall-clock gains (§IV-E), and ask whether the advantage generalizes across problem sizes (§IV-F). We then disentangle framework effects from algorithmic gains (§IV-G), examine what it costs FD in number of GPUs and energy to compensate AD speed (§IV-H), and characterize where time and memory concentrate inside the gradient pipeline (§IV-I). A. Experimental Setup TABLE I: Benchmark model characteristics. Pmin : minimum number of GPUs to fit the model. d

Model

Application

Latent

GST-S GST-M GST-L GST-T GST-C2 GST-C3

Synthetic Synthetic Synthetic Temperature Synthetic Synthetic

466 4 81K 4 1.0M 4 1.0M 4 8.5K 9 8.5K 15

SA1 AP1 WA1 WA2

Benchmark Air poll. Benchmark Benchmark

964K 606K 1.9M 646K

n Pmin

92 5 812 100 4,002 250 2,865 365 708 12 1,062 8

1 1 1 1 1 1

15 5,019 192 15 12,630 48 15 3,741 512 15 13,455 48

1 4 4 4

a) Benchmark Models: GST-S/M/L (Table I) are univariate synthetic benchmarks of increasing size; GST-T is a real-world temperature dataset [9] with daily observations at 2,865 weather stations over one year (nt =365). GST-C2 and GST-C3 add multivariate (coregionalization) structure with k=2 and k=3 variables, respectively [1]. AP1 is a real-world air pollution model that estimates particulate matter (PM2.5, PM10) and ozone (O3 ) concentrations over northern Italy [1]. SA1, WA1, and WA2 are production-scale benchmarks from DALIA [1] that exercise different scaling regimes: SA1 has moderate blocks over many time steps (nt =192), WA1 stresses temporal depth (nt =512), and WA2 stresses spatial density

AD-Dense AD-Loop AD-Ckpt

ADELIA

GST-S .34 / .4 .004 / <.1 .007 / <.1 .007 / <.1 .006 / <.1 GST-C2 1.6 / .2 .29 / 2.9 .075 / .9 .085 / 1.0 .079 / .2 GST-C3 2.9 / .2 .39 / 3.3 .083 / 1.3 .094 / 1.5 .096 / .3 GST-M 2.7 / 2.4 OOM .78 / 8.4 .87 / 6.4 .70 / 1.1 GST-L 56.6 / 76.9 OOM OOM OOM 23.8 / 63.3 GST-T 45.5 / 61.2 OOM OOM OOM 17.1 / 46.5

(ns =4,485). All four are trivariate models with d=15 hyperparameters controlling spatial range, temporal smoothness, and cross-variable coupling. b) Hardware: All experiments run on the CSCS Alps supercomputer [25] using NVIDIA GH200 Grace Hopper nodes, each with four Hopper-based GPUs (96 GiB HBM3 each) and a 72-core ARM Neoverse V2 CPU (Grace) with 480 GiB LPDDR5X memory, connected via NVLink-C2C. All four GPUs per node are utilized; multi-node runs communicate over an HPE Slingshot-11 interconnect. c) Software: We use ADELIA (extending DALIA [1]) with JAX 0.8.1 for AD and mpi4jax for distributed communication; DALIA with CuPy 13.6.0 provides baseline FD operations. MPI uses Cray MPICH with CPU-staged transfers (no CUDA-aware MPI or NCCL). We report absolute performance alongside speedups, average over at least 10 runs after warmup, and show 95% confidence intervals [26]. JIT compilation is a one-time cost (39–246 s across models), amortized over long optimization runs. B. Structure-Exploiting Differentiation Is the custom backward pass necessary, or could generic AD approaches work? We evaluate five strategies on a single GPU (Table II): FD (central finite differences, 2d+1 evaluations per gradient); ADDense (dense Cholesky on the full N ×N matrix, storing the dense factor); AD-Loop (BTA Cholesky as a lax.scan loop, storing all n loop carries and per-step L factors in the AD tape); AD-Ckpt (gradient checkpointing [27], trading recomputation for reduced tape storage); and ADELIA (storing only the n Schur carries and reconstructing L factors on the fly). FD scales to all sizes, but is the slowest strategy. AD-Dense OOMs at GST-M (81K latent variables); AD-Loop and ADCkpt OOM at GST-L and GST-T, where the stored states exceed GPU memory. Only ADELIA fits the million-variable models (GST-L within 63.3 GiB, GST-T within 46.5 GiB) while running 2.4× and 2.7× faster than FD, respectively. Structure-exploiting differentiation is thus a necessity, not merely an optimization: every generic AD strategy fails at increasing problem scale.

Objective f(θ)

Model

Temperature Monitoring (GST-T)

Gradient norm ‖∇f‖

TABLE II: Per-gradient time (s) / peak GPU memory (GiB) for five differentiation strategies. OOM = exceeded GPU memory. Best AD time in bold.

1.0B

ADELIA DALIA (FD)

500M 0 1.0B 500M

Air Pollution Monitoring (AP1) 145k

FD stalls on large models

142k 140k

Same convergence for small models

FD: ‖∇f‖ = 14k

20k 10k

AD: ‖∇f‖ = 1.4

0 5

Iteration

Fig. 3: L-BFGS convergence on two real-world models: objective (top) and gradient norm (bottom). Beyond the mathematical derivation in Section III, we validate correctness empirically and evaluate convergence quality. a) Exact reference: On models small enough for ADLoop (Table II), we use it as a reference since it relies on JAX’s built-in AD with no custom derivatives. ADELIA matches it: worst-case relative error 1.2 × 10−7 (GST-C3); smaller models to 10−12 or better. To validate the distributed twophase implementation, we additionally compare to the ADLoop reference on GST-C3 with P=2, achieving 6.8 × 10−6 . b) Finite-difference comparison: On large models, we validate against central finite differences (h = 10−3 ). All four distributed models achieve element-wise relative errors of 5 × 10−4 to 6 × 10−3 , consistent with FD truncation O(h2 ). 1) Convergence case study: Figure 3 compares L-BFGS convergence using ADELIA and DALIA from the same initial θ on two real-world models. The magnitude of the objective f varies across models; the goal is to minimize f , not for it to reach zero. On the smaller models (GST-S through GST-T; GST-T shown left), both methods converge to the same minimum. As model complexity grows, exact gradients become essential. On the largest distributed models the gap is significant: on AP1 (shown right), ADELIA converges in 165 iterations, reducing ∥∇f ∥ from 4,874 to 1.4; DALIA’s gradient norm is elevated because ∥∇f ∥ amplifies per-component FD errors across d=15 parameters, stalling after 19 iterations at f =145,660 vs. ADELIA’s f =139,618. On SA1, FD stalls at ∥∇f ∥=4,019 with a non-positive-definite Hessian, while ADELIA reaches ∥∇f ∥=1.5, a 2,700× reduction, with a well-conditioned Hessian for posterior uncertainty quantification. WA1 and WA2 follow suit, with DALIA stalling after 2–5 iterations and ADELIA reaching 104 × lower gradient norms. At this scale, ADELIA’s exact gradients are not just an optimization but a necessity: without them, the optimizer cannot reach a valid minimum. D. Minimum Resources Comparison

C. Gradient Correctness and Convergence Quality

With the minimum number of GPUs required to run a model, how much speedup does AD provide over FD?

Are the AD gradients numerically correct, and do they lead to high-quality convergence on real-world models?

We run each model on its minimum feasible GPU count (Pmin in Table I) and compare per-gradient times. Figure 4

Wall-clock time (s)

400 300

(a) Small/medium models Optimization (ADELIA) Hessian (ADELIA) Optimization (DALIA) Hessian (DALIA) Per-gradient speedup End-to-end speedup

200

30.4× 26.9×

3.9× 15.9×

100

20.0

50.9×

3.5×

17.5

Wall-clock time (h)

500

GST-S cAD = 0.2

4.6× 3.6×

10.0 7.5 5.0 2.5

GST-M cAD = 2.3

7.3×

12.5

12.9×

GST-C2 cAD = 1.2

7.9×

15.0

43.5×

(b) Large models Total of 26 hours saved across large models by using ADELIA

GST-C3 cAD = 1.0

0.0

6.4× 2.8×

2.4×

2.3×

2.2×

5.0×

4.2× 2.9×

GST-T GST-L WA1 AP1 WA2 SA1 cAD = 3.2 cAD = 3.7 cAD = 4.8 cAD = 7.4 cAD = 6.8 cAD = 3.9

Fig. 4: End-to-end wall-clock breakdown (optimization + Hessian) on Pmin GPUs, adjusted to the same number of iterations (using ADELIA as baseline) to isolate computational cost from convergence differences. presents the results: AD achieves per-gradient speedups of 2.4–50.9× across all ten models. The speedup is bounded by 2d+1 (the FD evaluation count), but the AD backward pass adds overhead, and per-evaluation framework efficiency play a role (Section IV-G). We capture this with the cost ratio cAD = TAD /teval , which measures how expensive one AD gradient is relative to one FD evaluation (teval = TFD /(2d+1)). The observed speedup is then (2d+1)/cAD : when cAD < 1, AD gains a per-evaluation bonus on top of the count reduction; when cAD > 1, it pays a penalty that partially offsets it. Three regimes emerge. a) Small, fully fused models (GST-S, GST-C2, GST-C3): XLA compiles the entire forward-backward computation into a single fused kernel, eliminating per-kernel launch overhead and bringing the full AD gradient close to or below the cost of a single CuPy-based FD evaluation (cAD ≲ 1), a framework advantage quantified in Section IV-G. b) Medium, compute-bound models (GST-M, GST-L, GST-T, SA1): The backward pass costs roughly 1.3–2.4× the forward, placing the total AD gradient at 2.3–3.4× a single forward evaluation, well within the theoretical 3–5× bound on the backward-to-forward ratio [10]. c) Distributed models (AP1, WA1, WA2): These models require 4 GPUs to fit in memory; Per-block JIT dispatch, CPU↔GPU transfers for CPU-staged carries, and the twophase reduced system overhead add framework cost, raising cAD to 4.8–7.4. Despite this penalty, the 2d+1=31 evaluation count advantage still yields 4.2–6.4× speedups. E. End-to-End Wall-Clock How do per-gradient speedups translate to end-to-end wallclock time? Figure 4 decomposes the INLA run into optimization and posterior Hessian computation. As shown in Section IV-C, FD and AD can follow different optimization trajectories. To isolate computational cost from convergence quality, our endto-end comparison uses a fixed-work budget: we run ADELIA

until convergence for each model, and use the resulting iteration count for both methods. The Hessian stage (Section III-D) adds a second advantage: AD computes each Hessian column via two gradient evaluations, while FD requires O(d2 ) objective function evaluations via second-order finite differences, yielding 1.1–23.3× Hessian speedups. Overall end-to-end speedups range from 2.2× (GST-L) to 7.3× (SA1) for the production-scale models, with distributed models at 2.9–5.0×. In absolute terms, AD reduces the projected INLA run from 14.6 h to 2.0 h for SA1 and from 9.6 h to 2.7 h for WA2 on equal hardware. F. Problem Size Impact Study Does AD’s advantage hold consistently across problem sizes, or is it model-specific? We study how speedup scales with problem size along two axes using the WA1 and WA2 model families on a single GPU. The largest sizes here are smaller than the production configurations in Table I, which require multiple GPUs; this study isolates the scaling trend from distribution overhead. For WA1 temporal scaling (Figure 5a), cAD rises from 1.4 to 3.7, converging quickly. At small n, the backward pass costs 1.6× the forward; at large n, the loop is compute-bound (β = 2.2), pushing cAD to its asymptote. Figure 5(b) shows the same pattern for WA2 spatial scaling (ns =72 to 1,119): cAD rises from 1.2 to 3.5 as O(b3 ) computation overtakes kernel-launch overhead. Both axes converge to cAD ≈ 3.5, confirming that the asymptotic cost ratio is independent of whether the problem grows in time or space. AD retains 8–25× speedups across the full range. G. Framework Effect Decomposition How much of the observed speedup comes from XLA compilation efficiency versus the algorithmic advantage of AD? ADELIA uses JAX/XLA, while DALIA’s FD baseline uses CuPy with Serinv [16]. We isolate the framework effect by

8.4×

100 0

8.3×

2 cAD =1.4

8.5×

8.9×

21.8×

32 cAD =3.5

64 cAD =3.6

128 cAD =3.7

256 cAD =3.7

Number of time steps (nt )

Per-gradient time (s)

(a) WA1: Temporal scaling

(b) WA2: Spatial scaling 20

ADELIA DALIA 13.6×

24.8×

72 cAD =1.2

8.9×

282 cAD =2.3

Spatial mesh nodes (ns )

1119 cAD =3.5

Fig. 5: AD speedup vs. problem size. (a) WA1 with increasing temporal resolution (nt = 2 to 256, ns =1,247). (b) WA2 with increasing spatial resolution (ns = 72 to 1,119, nt =48). Speedup annotations on FD bars; cAD values below each size label. 4.2×

GST-S r=12.4

GST-M r=1.0 GST-L r=0.9 GST-T r=0.9 SA1 r=0.9

45.7×

7.6×

GST-C2 r=2.2 GST-C3 r=2.8

Observed vs. algorithmic speedup

10.8× 3.9× 4.0× 2.4× 2.8× 2.8× 3.0× 7.9× 4.3×

AP1 r=0.1

38.0×

Algorithmic (r = 1) Observed

38.6×

r<1 CuPy faster

27.4×

4.4×

WA2 r=0.1

r≈1 parity

9.1×

6.3×

WA1 r=0.2

r>1 JAX faster

17.7×

34.0×

Speedup over FD

Fig. 6: Framework effect decomposition: observed speedup (pink) vs. algorithmic speedup at framework parity (gold). running the same evaluation on both stacks, measuring r = tSerinv /tJAX . At framework parity (r=1), the speedup would be purely algorithmic: (2d+1)/(1+β), where β = tbwd /tfwd is the backward-to-forward cost ratio. Figure 6 decomposes each model’s speedup into algorithmic and framework contributions. For small blocks (GST-S), XLA kernel fusion yields r ≫ 1, amplifying the algorithmic speedup. GST-C2 and GSTC3 also benefit from r > 1 (r=2.2 and 2.8); their short temporal loops (n=8–12) favor XLA’s compiled scan over Serinv’s per-iteration Python dispatch. For compute-bound models (GST-M, GST-L, GST-T, SA1), both stacks dispatch to cuBLAS/cuSOLVER and r ≈ 0.9–1.0, close to parity. For distributed models, r drops well below parity (0.1–0.25): Serinv overlaps CPU↔GPU transfers with computation via direct CUDA stream control, whereas JAX’s per-block JIT dispatch is synchronous. This is a framework-level engineering limitation, not a fundamental cost of AD.

We measure per-gradient time for FD on increasing GPU counts. FD scales by running its 2d+1 evaluations in parallel: given P GPUs, it runs ⌊P/Pmin ⌋ concurrent evaluations, each on Pmin GPUs. Since AD replaces these 2d+1 evaluations with a single forward-backward pass, it is run on its minimum configuration only. Energy is measured via Cray PM hardware counters at the node power-supply level (10 Hz, capturing GPU, CPU, memory, and interconnect), with MPI barriers ensuring the measurement window includes idle time to reflect the real allocation cost. a) Per-gradient time scaling: Figure 7(a) focuses on the four production-scale models, all with d=15 (2d+1=31 FD evaluations). Under ideal parallelism, FD matches AD at P ∗ ≈ Pmin · (2d+1)/cAD GPUs. SA1 (cAD =3.9, single-GPU AD, Pmin =1) starts at 7.9× with one FD GPU and crosses below 1 between 8 and 16 GPUs; its curve stops at 32 GPUs because 31 ≤ 32/Pmin , so all evaluations already run in a single round and additional GPUs sit idle. The distributed models (AP1, WA1, WA2; cAD =4.8–7.4, 4-GPU AD) cross below 1 between 16 and 32 FD GPUs and are scaled to 128 GPUs. For the smaller benchmarks (not shown in the figure), the pattern is similar: low-d models (d=4) break even at 4–8 GPUs, while high-d models (d=9–15) remain above break-even through 32 GPUs. b) Energy efficiency: Figure 7(b) shows the opposite trend: unlike time ratios, which decrease as FD scales to more GPUs, energy ratios are flat or increasing. Under ideal parallelism, doubling GPUs halves FD’s wall time but doubles the hardware, so total energy ( time × per-GPU power × GPUs) is conserved: parallelism trades time for GPUs, leaving total work unchanged. In practice, sub-linear scaling and perGPU base power draw makes the total energy grow with GPU count even as time decreases. The best case for FD (WA1 at 128 GPUs, where FD is ∼4.2× faster), it still consumes 7.3× more energy. Crucially, parallelism offers no remedy for gradient quality: The convergence gap shown in Section IV-C persists regardless of scale. ADELIA is therefore the only path to both energy-efficient and high-quality gradient computation.

H. Resource and Energy Efficiency

I. Performance Analysis

FD’s 2d+1 evaluations are embarrassingly parallel. How many GPUs does FD need to match AD’s speed, and at what energy cost?

Where do the time and memory costs concentrate in practice, and what are the practical deployment overheads?

ADELIA faster

6 4 2 DALIA faster

SA1

AP1

128

DALIA energy / ADELIA energy

DALIA time / ADELIA time

(a) Per-gradient time ratio 8

(b) Measured energy ratio 8

ADELIA cheaper

6 4 2 0

DALIA cheaper 1

Number of DALIA GPUs

WA1

WA2

128

Per-gradient time (s)

Fig. 7: Resource and energy efficiency for the four production-scale models (d=15). (a) Per-gradient speedup as FD scales to more GPUs; AD runs on Pmin . (b) Energy ratio. Values above dashed lines favor AD. 150 100 50 0

Chol + L-TriSolv U-TriSolv + SI Prior grad Quad grad

55%

52%

51%

WA1

AP1

WA2

Fig. 8: Per-stage time breakdown for distributed models on 4 GPUs. Percentages show the combined Python-loop fraction. a) Peak GPU memory: ADELIA’s peak GPU memory is dominated by the n stored Schur complement carries (b × b each). Total memory ranges from < 1 MiB (GSTS) to 63.3 GiB (GST-L) for single-GPU models (Table II); GST-T uses 46.5 GiB. For distributed models, the CPU-staging strategy from Section III-C keeps per-GPU memory to a single block’s workspace, with Schur carries stored in CPU memory. b) Stage-level profiling: Figure 8 breaks down the pergradient time for the three distributed models into the four gradient stages from Section III-A. The two Python-loop stages (Chol + L-TriSolv, U-TriSolv + SI) account for 51– 55% of total time; the prior log-determinant gradient (Phase C) accounts for the remainder 44–47%, scaling with block size. c) Two-phase execution profile: Phase 1 (Section III-C) dominates total time. Phase 2 overhead varies by model: negligible for WA1 (∼7%, b=3,741) but ∼24% for AP1/WA2 (b>12,000), driven by the reduced system’s allgather of O(P · b2 ) elements (22–25 s) and factorization (6–8 s). V. R ELATED W ORK a) INLA Implementations: R-INLA [7], [28] is the standard INLA implementation, using FD for gradients and PARDISO [29] for Cholesky and SI. DALIA [1] scales INLA to hundreds of GPUs by parallelizing across FD evaluations, precision matrix assembly, and distributed BTA solvers from Serinv [16], but retains finite differences. The gmrfs library [30] uses JAX AD for INLA on GMRFs but differentiates through generic sparse Cholesky via Python loops without a custom backward pass, analogous to but less efficient than our ADLoop baseline (Table II), which uses compiled lax.scan.

ADELIA replaces FD in DALIA with structure-exploiting AD while retaining the same solver infrastructure. b) Sparse Solvers and AD for Structured Matrices: General-purpose sparse direct solvers (PARDISO [29], MUMPS [31], SuperLU DIST [32], cuSOLVER) use fillreducing orderings that disrupt block-tridiagonal structure. Serinv [16] operates on the BTA structure for O(n · b3 ) factorization; Our carry-based reconstruction can be seen as a structure-aware form of gradient checkpointing [27], exploiting block sparsity by storing Schur complement carries and reconstructing factors on the fly. Durrande et al. [33] derived reverse-mode AD rules for banded Cholesky and Takahashi’s equations on scalar banded matrices, but their approach does not handle block structure, arrowhead coupling, or GPU execution. Algorithmically, the backward pass relates to the Rauch–Tung–Striebel smoother [34], but differentiable state-space models target dense states of tens to hundreds, whereas our spatial blocks reach b=13,455. c) Probabilistic Programming and GPU Inference: AD-based frameworks such as Stan [35], PyMC [36], NumPyro [37], and TensorFlow Probability [38] provide general Bayesian inference but do not exploit BTA sparsity; naive generic sparse AD destroys block structure. TMB [12] combines tape-based AD (CppAD) with CHOLMOD’s sparse Cholesky and the inverse subset algorithm to differentiate the Laplace-approximated marginal likelihood, supporting automatic sparsity detection for models with up to 106 random effects. However, TMB uses general-purpose fill-reducing orderings that do not preserve BTA block structure, precluding the O(n · b3 ) block-sequential algorithms central to our approach, and is limited to CPU execution. Margossian et al. [13], [14] derive an adjoint-differentiated Laplace approximation that, like our approach, exploits the envelope theorem to avoid implicit differentiation through the mode; their 2023 generalization [14] compares the inverse-subset and adjoint gradient strategies. The Stan implementation does not exploit structured sparsity and lacks GPU support, limiting scalability to the models supported by ADELIA. None target INLA’s combination of Laplace approximation, structured sparse factorization, and SI.

VI. C ONCLUSION We presented ADELIA, the first AD-enabled INLA framework that exploits the block sparsity of the underlying precision matrices to compute exact gradients on multiple GPUs, scaling to models with up to 1.9M latent variables. On four production-grade multivariate models, AD achieves 4.2– 7.9× per-gradient speedups on equal hardware, with 5–8× energy savings even when FD matches wall-clock time by parallelizing over more nodes. Beyond performance, exact gradients are essential for the largest multivariate models, where FD stalls at gradient norms orders of magnitude above tolerance. Regardless of resource availability, ADELIA is the go-to option for fast, energy-efficient, accurate gradients for HPC-scale INLA applications. Enabling reliable convergence at these scales, at a cost independent of the number of hyperparameters, makes it possible to ask scientific questions that the FD computational overhead and approximation error previously prevented. ACKNOWLEDGMENTS All sections of this paper were iteratively refined using Claude by Anthropic for grammar and phrasing improvements. R EFERENCES [1] L. Gaedke-Merzhäuser, V. Maillou, F. R. Avellaneda, O. Schenk, M. Luisier, P. Moraga, A. N. Ziogas, and H. Rue, “Accelerated spatio-temporal bayesian modeling for multivariate gaussian processes,” 2025. [Online]. Available: https://doi.org/10.48550/arXiv.2507.06938 [2] A. Gelman, J. B. Carlin, H. S. Stern, D. B. Dunson, A. Vehtari, and D. B. Rubin, Bayesian Data Analysis, 3rd ed. Chapman and Hall/CRC, 2013. [Online]. Available: https://doi.org/10.1201/b16018 [3] C. K. Wikle, A. Zammit-Mangion, and N. Cressie, Spatio-temporal statistics with R, 1st ed. Chapman and Hall/CRC, 2019. [Online]. Available: https://doi.org/10.1201/9781351769723 [4] M. Blangiardo and M. Cameletti, Spatial and spatio-temporal Bayesian models with R-INLA. John Wiley & Sons, Ltd, 2015. [Online]. Available: https://doi.org/10.1002/9781118950203 [5] P. Moraga, Geospatial Health Data: Modeling and visualization with R-INLA and Shiny, 1st ed. Chapman and Hall/CRC, 2019. [Online]. Available: https://doi.org/10.1201/9780429341823 [6] S. Brooks, A. Gelman, G. Jones, and X.-L. Meng, Handbook of Markov Chain Monte Carlo, 1st ed. Chapman and Hall/CRC, 2011. [Online]. Available: https://doi.org/10.1201/b10905 [7] H. Rue, S. Martino, and N. Chopin, “Approximate bayesian inference for latent gaussian models by using integrated nested laplace approximations,” Journal of the Royal Statistical Society Series B: Statistical Methodology, vol. 71, no. 2, pp. 319–392, 04 2009. [Online]. Available: https://doi.org/10.1111/j.1467-9868.2008.00700.x [8] F. Lindgren, H. Bakka, D. Bolin, E. Krainski, and H. Rue, “A diffusion-based spatio-temporal extension of gaussian matérn fields,” 2023. [Online]. Available: https://doi.org/10.48550/arXiv.2006.04917 [9] L. Gaedke-Merzhäuser, E. Krainski, R. Janalik, H. Rue, and O. Schenk, “Integrated nested laplace approximations for largescale spatiotemporal bayesian modeling,” SIAM Journal on Scientific Computing, vol. 46, no. 4, pp. B448–B473, 2024. [Online]. Available: https://doi.org/10.1137/23M1561531 [10] A. Griewank and A. Walther, Evaluating Derivatives, 2nd ed. Society for Industrial and Applied Mathematics, 2008. [Online]. Available: https://doi.org/10.1137/1.9780898717761 [11] A. G. Baydin, B. A. Pearlmutter, A. A. Radul, and J. M. Siskind, “Automatic differentiation in machine learning: a survey,” 2018. [Online]. Available: https://doi.org/10.48550/arXiv.1502.05767 [12] K. Kristensen, A. Nielsen, C. W. Berg, H. Skaug, and B. M. Bell, “TMB: Automatic differentiation and Laplace approximation,” Journal of Statistical Software, vol. 70, no. 5, p. 1–21, 2016. [Online]. Available: https://doi.org/10.18637/jss.v070.i05

[13] C. C. Margossian, A. Vehtari, D. Simpson, and R. Agrawal, “Hamiltonian Monte Carlo using an adjoint-differentiated Laplace approximation: Bayesian inference for latent Gaussian models and beyond,” 2020. [Online]. Available: https://doi.org/10.48550/arXiv.2004. 12550 [14] C. C. Margossian, “General adjoint-differentiated Laplace approximation,” 2023. [Online]. Available: https://doi.org/10.48550/ arXiv.2306.14976 [15] H. Rue and L. Held, Gaussian Markov Random Fields: Theory and Applications. Chapman and Hall/CRC, 2005. [Online]. Available: https://doi.org/10.1201/9780203492024 [16] V. Maillou, L. Gaedke-Merzhäuser, A. N. Ziogas, O. Schenk, and M. Luisier, “Serinv: A scalable library for the selected inversion of block-tridiagonal with arrowhead matrices,” 2025. [Online]. Available: https://doi.org/10.48550/arXiv.2503.17528 [17] F. Lindgren, H. Rue, and J. Lindström, “An explicit link between Gaussian fields and Gaussian Markov random fields: the stochastic partial differential equation approach,” Journal of the Royal Statistical Society Series B: Statistical Methodology, vol. 73, no. 4, pp. 423–498, 08 2011. [Online]. Available: https://doi.org/10.1111/j.14679868.2011.00777.x [18] J. Nocedal and S. J. Wright, Numerical Optimization, 2nd ed. Springer, 2006. [Online]. Available: https://doi.org/10.1007/978-0-387-40065-5 [19] I. Murray, “Differentiation of the Cholesky decomposition,” 2016. [Online]. Available: https://doi.org/10.48550/arXiv.1602.07527 [20] P. Milgrom and I. Segal, “Envelope theorems for arbitrary choice sets,” Econometrica, vol. 70, no. 2, pp. 583–601, 2002. [Online]. Available: https://doi.org/10.1111/1468-0262.00296 [21] K. B. Petersen and M. S. Pedersen, “The matrix cookbook,” 2012, version 20121115. [Online]. Available: http://www2.compute.dtu.dk/ pubdb/pubs/3274-full.html [22] A. M. Schmidt and A. E. Gelfand, “A bayesian coregionalization approach for multivariate pollutant data,” Journal of Geophysical Research: Atmospheres, vol. 108, no. D24, 2003. [Online]. Available: https://doi.org/10.1029/2002JD002905 [23] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, Y. Katariya, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang, “JAX: composable transformations of Python+NumPy programs,” 2018. [Online]. Available: http://github. com/jax-ml/jax [24] D. Häfner and F. Vicentini, “mpi4jax: Zero-copy MPI communication of JAX arrays,” Journal of Open Source Software, vol. 6, no. 65, p. 3419, 2021. [Online]. Available: https://doi.org/10.21105/joss.03419 [25] M. Martinasso, M. Klein, and T. Schulthess, “Alps, a versatile research infrastructure,” in Proceedings of the Cray User Group, ser. CUG ’25. New York, NY, USA: Association for Computing Machinery, 2025, pp. 156–165. [Online]. Available: https://doi.org/10.1145/3757348.3757365 [26] T. Hoefler and R. Belli, “Scientific benchmarking of parallel computing systems: Twelve ways to tell the masses when reporting performance results,” in Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC), ser. SC ’15. Association for Computing Machinery, 2015. [Online]. Available: https://doi.org/10.1145/2807591.2807644 [27] T. Chen, B. Xu, C. Zhang, and C. Guestrin, “Training deep nets with sublinear memory cost,” 2016. [Online]. Available: https://doi.org/10.48550/arXiv.1604.06174 [28] L. Gaedke-Merzhäuser, J. van Niekerk, O. Schenk, and H. Rue, “Parallelized integrated nested Laplace approximations for fast Bayesian inference,” Statistics and Computing, vol. 33, p. 25, 2023. [Online]. Available: https://doi.org/10.1007/s11222-022-10192-1 [29] O. Schenk and K. Gärtner, “Solving unsymmetric sparse systems of linear equations with PARDISO,” Future Generation Computer Systems, vol. 20, no. 3, pp. 475–487, 2004, selected numerical algorithms. [Online]. Available: https://doi.org/10.1016/j.future.2003.07.011 [30] A. Geraschenko, “gmrfs: INLA for Gaussian Markov random fields in JAX,” 2024. [Online]. Available: https://github.com/geraschenko/gmrfs [31] P. R. Amestoy, I. S. Duff, J.-Y. L’Excellent, and J. Koster, “A fully asynchronous multifrontal solver using distributed dynamic scheduling,” SIAM Journal on Matrix Analysis and Applications, vol. 23, no. 1, pp. 15–41, 2001. [Online]. Available: https: //doi.org/10.1137/S0895479899358194 [32] X. S. Li and J. W. Demmel, “SuperLU DIST: A scalable distributedmemory sparse direct solver for unsymmetric linear systems,” ACM

Trans. Math. Softw., vol. 29, no. 2, pp. 110–140, 2003. [Online]. Available: https://doi.org/10.1145/779359.779361 [33] N. Durrande, V. Adam, L. Bordeaux, S. Eleftheriadis, and J. Hensman, “Banded matrix operators for Gaussian Markov models in the automatic differentiation era,” 2019. [Online]. Available: https: //doi.org/10.48550/arXiv.1902.10078 [34] H. E. RAUCH, F. TUNG, and C. T. STRIEBEL, “Maximum likelihood estimates of linear dynamic systems,” AIAA Journal, vol. 3, no. 8, pp. 1445–1450, 1965. [Online]. Available: https://doi.org/10.2514/3.3166 [35] B. Carpenter, A. Gelman, M. D. Hoffman, D. Lee, B. Goodrich, M. Betancourt, M. Brubaker, J. Guo, P. Li, and A. Riddell, “Stan: A probabilistic programming language,” Journal of Statistical Software, vol. 76, no. 1, pp. 1–32, 2017. [Online]. Available: https://doi.org/10.18637/jss.v076.i01 [36] O. Abril-Pla, V. Andreani, C. Carroll, L. Dong, C. J. Fonnesbeck, M. Kochurov, R. Kumar, J. Lao, C. C. Luhmann, O. A. Martin, M. Osthege, R. Vieira, T. Wiecki, and R. Zinkov, “PyMC: a modern, and comprehensive probabilistic programming framework in Python,” PeerJ Computer Science, vol. 9, p. e1516, 2023. [Online]. Available: https://doi.org/10.7717/peerj-cs.1516 [37] D. Phan, N. Pradhan, and M. Jankowiak, “Composable effects for flexible and accelerated probabilistic programming in NumPyro,” in NeurIPS Workshop on Program Transformations, 2019. [Online]. Available: https://doi.org/10.48550/arXiv.1912.11554 [38] J. V. Dillon, I. Langmore, D. Tran, E. Brevdo, S. Vasudevan, D. Moore, B. Patton, A. Alemi, M. Hoffman, and R. A. Saurous, “TensorFlow distributions,” 2017. [Online]. Available: https://doi.org/10.48550/arXiv. 1711.10604

Record · ID 168263 · SHA-256 e1b07ccd2fea851d
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.