Variational Inference Using a Differentiable Multigrid Linear Solver

2026-08-01

AI Summary

This work addresses the computational challenge of gradient evaluation in high-dimensional forward models arising in PDE-constrained inverse problems by introducing a Differentiable Multigrid Solver (DMGS). For the first time, DMGS enables compatibility between exact adjoint operations and automatic differentiation frameworks. The method explicitly derives the full multigrid hierarchy’s adjoint operators for steady-state diffusion-absorption problems and integrates them into JAX, efficiently supporting both Jacobian-vector and vector-Jacobian products. Applied to a 3D tissue diffuse radiative transfer inverse problem, DMGS achieves accurate reconstruction of effective radiation sources (χ² = 1.1) with reduced peak memory usage and controllable backward-pass overhead, while demonstrating strong generalization across 32 validation cases.

Abstract

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.

Problem

• variational inference
• multigrid solver
• automatic differentiation
• PDE-constrained inverse problems
• adjoint operations

Innovation

• differentiable multigrid
• variational inference
• adjoint solver
• automatic differentiation
• PDE-constrained inverse problems