基于有限元的可微物理中的隐式微分（涉及二阶导数）及基准测试

《Computer Physics Communications》：Implicit differentiation with second-order derivatives and benchmarks in finite-element-based differentiable physics

本研究提出了一种统一框架，用于在有限元方法（FEM）中实现隐式二阶微分，特别是Hessian矩阵的计算。通过JAX的自动微分工具，验证了Hessian-向量积的准确性，并通过四个基准测试（线性/非线性、2D/3D、单变量/耦合变量）展示了在非线性逆问题（如牵引力识别、形状优化）中牛顿-CG方法利用精确Hessian加速收敛的效果。该工作为将二阶隐式微分整合到可微分物理引擎奠定了基础，促进PDE约束优化问题的更快速和可靠的求解。

　　

Tianju Xue

香港科技大学土木与环境工程系，中国香港

摘要

微分编程通过实现数值模拟的自动微分（AD），正在彻底改变计算科学。虽然一阶梯度已经得到广泛应用，但在基于有限元的微分物理中，隐函数的二阶导数（Hessian）仍然研究不足。本研究通过推导并实现一个用于偏微分方程（PDE）约束优化问题中隐函数Hessian计算的框架，填补了这一空白。我们利用基本的AD工具（雅可比矩阵-向量积/向量-雅可比矩阵积）构建了一个Hessian-向量积算法，并通过有限差分近似验证了其准确性。四个涵盖线性/非线性、2D/3D以及单变量/多变量问题的基准测试展示了二阶信息的实用性。结果表明，使用精确Hessian的牛顿-CG方法可以加速非线性逆问题的收敛速度（例如牵引力识别、形状优化），而对于线性问题，L-BFGS-B方法已经足够。我们的工作为将二阶隐式微分集成到微分物理引擎中奠定了基础，从而实现了更快、更可靠的PDE约束优化。

项目概述

项目名称： CPC库链接到项目文件： https://github.com/tianjuxue/hessian Code Ocean胶囊：许可条款： GPLv3 编程语言： Python 补充材料： 无 问题性质：基于有限元的微分物理缺乏隐函数的Hessian计算，这阻碍了优化解的快速收敛。 解决方法：采用离散化-优化框架，通过自动微分原语来计算隐函数Hessian-向量积。

引言

微分编程的出现——一种能够实现整个数值模拟自动微分（AD）的范式转变——正在彻底改变计算科学。与传统符号或数值微分不同，微分编程将AD直接嵌入到数值求解器的计算图中，从而无需手动推导即可精确计算梯度[1]。

微分编程的典型应用是在深度学习领域，其中更为人熟知的“反向传播”（即AD的反向模式）一直是训练深度神经网络的引擎[2]。在计算科学领域，微分编程正在迅速发展，因为可以将微分编程应用于经典的基于物理的数值求解器，从而获得易于进行逆优化的微分物理引擎。如图1所示，深度神经网络训练与微分物理之间的并排比较揭示了它们之间的许多相似之处：两种工作流程都有输入x和输出y；前向函数（无论是神经网络还是数值求解器）由θ参数化；存在损失函数l(y, θ)（或目标函数g(y, θ)；目标是通过基于梯度的训练（或优化）算法找到所需的θ。

像JAX[3]或PyTorch[4]这样的框架使微分编程普及化，促进了微分物理在许多领域的应用，如流体力学[5]、[6]、分子动力学[7]、固体力学[8]、[9]、热传递[10]等。然而，在深度学习中寻找梯度与在微分物理中寻找梯度之间存在一个关键技术差异：在深度学习中，神经网络函数通常是显式的，而物理系统的数值求解器通常是隐式的。直接后果是，作为θ和x函数的输出y通常是显式的，因此可以直接应用AD；而在微分物理中，这种关系通常是隐式的，由于求解非线性方程的内部迭代，无法直接应用简单的AD。大多数关于微分物理的先前工作都采用了显式的数值求解器来绕过这个问题。处理与隐函数相关的梯度的正确方法是将伴随方法[11]集成到定制的AD规则中，如[12]中所详细讨论的。

基于伴随方法支持的定制AD规则，我们在之前的工作JAX-FEM[13]中提出了一种可处理隐函数的微分有限元方法（FEM）。我们选择有限元方法[14]作为离散化物理系统的工具，因为它是求解偏微分方程（PDEs）的最强大方法之一。该方法原则上也适用于其他数值方法，如有限体积方法[15]或无网格方法[16]。在微分物理的范围内，JAX-FEM有效地处理了数学上表述为PDE约束优化问题的逆问题[17]。

尽管基于伴随方法的隐式微分编程在一阶微分物理方面取得了最新进展，但二阶隐式微分（即Hessian矩阵）仍然研究不足。Hessian矩阵编码了目标函数景观的曲率信息，使牛顿型算法（例如牛顿-CG [18]、[19]）能够在最小值附近实现二次收敛率[20]、[21]。利用Hessian信息加速优化收敛的潜在优势在微分编程社区中非常吸引人（参见隐式微分包JAXopt [12]讨论论坛中的问题#474），但目前还无法获得用于隐式微分的Hessian信息。一般来说，解决PDE约束优化问题有两种策略：先优化后离散化[22]和先离散化后优化[23]。现有的确定Hessian信息的努力大多遵循先优化后离散化的策略[24]，例如用于贝叶斯逆问题的hIPPYlib [25]包。然而，微分物理的工作流程倾向于先离散化后优化策略，因为它提供了更高的灵活性，并与AD更深入地集成，展现了更高级别的自动化。

本研究通过开发一个统一且实用的框架，填补了经典基于伴随方法的Hessian方法与现代微分编程之间的空白，用于基于FEM的微分物理中的二阶隐式微分。本工作的关键特点包括：

1.

推导：我们遵循先离散化后优化的策略，在隐式微分编程范式下推导Hessian矩阵。我们的方法可以视为Blondel等人[12]的一阶框架的二阶推广。

2.

实现：我们使用JAX的原语函数雅可比矩阵-向量积（JVP）和向量-雅可比矩阵积（VJP）来实现Hessian-向量积，并通过有限差分近似方法验证其准确性并进行效率比较。

3.

基准测试：我们设计了四个基准逆问题（线性/非线性、2D/3D、单变量/多变量），涵盖了几个典型任务：参数识别、变形控制和形状优化。根据结果，我们了解了何时二阶信息对优化收敛更有帮助。

本文的结构如下。第2节定义了隐式微分的概念，提出了一个模型问题，并给出了Hessian矩阵的推导方法。第3节展示了隐式Hessian-向量积实现的细节。第4节通过四个代表性的数值基准测试讨论了使用Hessian信息帮助优化收敛的有效性。第5节进行了总结。

我们的代码可在https://github.com/CMSL-HKUST/hessian获取。

章节片段

隐式微分

在本节中，我们首先介绍了微分编程背景下的隐式微分概念。以非线性泊松问题为例，作为PDE约束优化的模型问题，其中离散化使用有限元方法完成。借助拉格朗日函数，我们推导了目标函数关于优化的一阶导数（梯度）和二阶导数（Hessian）的表达式。

实现

在本节中，我们继续推导Hessian信息，考虑了Hessian-向量积的实现。我们展示了Hessian-向量积的实现方式，并介绍了三种AD模式组合方式。通过与其他方法的比较以及泰勒余项测试验证了准确性，并讨论了这三种实现方式的效率。

数值示例

为了解决PDE约束优化问题，我们探讨了两种流行的优化算法：L-BFGS-B方法[31]和牛顿-CG方法[18]、[19]。这两种方法都依赖于导数的信息。然而，L-BFGS-B方法通过梯度历史近似来避免显式的Hessian-向量积，而牛顿-CG方法则直接依赖Hessian-向量积来迭代求解牛顿方程。

结论

我们的工作推进了将二阶隐式微分集成到基于有限元的微分物理中，解决了PDE约束优化中的一个紧迫需求。我们提出了一个在先离散化后优化范式下计算Hessian的框架，将伴随方法推广到二阶导数。通过利用JAX的自动微分原语（JVP/VJP），我们实现了Hessian-向量积，并通过有限差分方法验证了其准确性。

CRediT作者贡献声明

Tianju Xue：撰写——审阅与编辑、撰写——原始草稿、可视化、验证、软件、方法论、资金获取、形式分析、概念化。

利益冲突声明

作者声明他们没有已知的竞争财务利益或个人关系可能影响本文所述的工作。