2026/08/01 by Andrés Ramírez, Philipp Haim, David Föger +9
Mathematics · Computer Science · Physics and Astronomy · #math.NA #cs.NA #physics.med-ph #stat.CO #msc:65N55 #msc:65N06 #msc:62F15 #acm:65N55 #acm:65N06 #acm:62F15
18 pages, 13 figures, 1 table
arxiv created 2026/08/01 · arxiv updated 2026/08/04
Gradient-based Bayesian inference methods require efficient access to Jacobian and adjoint-Jacobian operators of high-dimensional forward models. While multigrid solvers provide near-optimal complexity for elliptic partial differential equations, they are rarely available in forms compatible with automatic differentiation (AD). We develop a differentiable multigrid solver for steady-state diffusion-absorption problems and derive its adjoint operations analytically through the full multigrid hierarchy. The resulting solver, DMGS, is implemented in C++ and interfaced with JAX to provide efficient Jacobian-vector and vector-Jacobian products for variational inference in the NIFTy framework. We validate the approach on a 3D inverse problem involving diffuse radiative transfer in tissue, reconstructing an effective radiative source from Monte Carlo-simulated data. The reconstruction reproduces the data at a reduced chi-squared of 1.1 and generalizes to 32 independent validation datasets. Benchmarks against a JAX-native multigrid implementation show comparable runtimes and consistently lower peak memory for the hand-derived adjoint, with modest reverse-mode overhead. These results establish differentiable multigrid solvers as practical building blocks for variational inference in PDE-constrained problems.