使用可微多重网格线性求解器的变分推断
Variational Inference Using a Differentiable Multigrid Linear Solver
浏览论文内容
中文总结 AI 辅助
该研究开发了可微多重网格求解器DMGS，将其与JAX对接用于NIFTy框架的变分推断，在3D组织辐射传输逆问题中验证了其高效性与泛化能力。
中文摘要 AI 辅助
基于梯度的贝叶斯推断方法需要高效获取高维前向模型的雅可比和伴随雅可比算子。多重网格求解器对椭圆偏微分方程（PDE）具有近最优复杂度，但很少有兼容自动微分（AD）的形式。我们开发了针对稳态扩散-吸收问题的可微多重网格求解器，并通过完整多重网格层次解析推导其伴随操作。所得求解器DMGS用C++实现，与JAX对接，为NIFTy框架中的变分推断提供高效的雅可比-向量和向量-雅可比乘积。我们在涉及组织中漫射辐射传输的3D逆问题上验证该方法，从蒙特卡洛模拟数据中重建有效辐射源。该重建以1.1的约化卡方值复现数据，并能泛化到32个独立验证数据集。与JAX原生多重网格实现的基准测试显示，手工推导的伴随操作具有相当的运行时间和始终更低的峰值内存，仅存在适度的反向模式开销。这些结果确立了可微多重网格求解器作为PDE约束问题变分推断的实用构建模块。
英文摘要
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.
