[论文解读] $ξ$-torch: differentiable scientific computing library

Muhammad Kasim, S. M. Vinko|arXiv (Cornell University)|Oct 5, 2020

Model Reduction and Neural Networks参考文献 14被引用 4

一句话总结

本文介绍了 $\xi$-torch，一个基于 PyTorch 的库，提供用于科学计算的可微分函数式，如根求解器、常微分方程（ODE）求解器和积分方法，并采用解析导数以提升数值稳定性与内存效率。该库通过支持一阶及更高阶导数，实现了端到端可微分的物理信息机器学习，在镜面设计与分子动力学优化中实现了快速收敛。

ABSTRACT

Physics-informed learning has shown to have a better generalization than learning without physical priors. However, training physics-informed deep neural networks requires some aspect of physical simulations to be written in a differentiable manner. Unfortunately, some operations and functionals commonly used in physical simulations are scattered, hard to integrate, and lack higher order derivatives which are needed in physical simulations. In this work, we present $ξ$-torch, a library of differentiable functionals for scientific simulations. Example functionals are a root finder and an initial value problem solver, among others. The gradient of functionals in $ξ$-torch are written based on their analytical expression to improve numerical stability and reduce memory requirements. $ξ$-torch also provides second and higher order derivatives of the functionals which are rarely available in existing packages. We show two applications of this library in optimizing parameters in physics simulations. The library and all test cases in this work can be found at https://github.com/xitorch/xitorch/ and the documentation at https://xitorch.readthedocs.io.

研究动机与目标

• 为解决现有库在物理信息深度学习中缺乏集成化、可微分科学函数式的问题。
• 通过使用解析导数而非前向算法的自动微分，提升科学模拟中的数值稳定性并降低内存使用。
• 为高级科学模拟（如振动频率计算）提供高阶导数（例如二阶导数）的支持。
• 通过提供统一、模块化的库，结合纯函数与线性算子抽象，促进深度学习与物理模拟的集成。
• 通过支持通过复杂模拟组件（如 ODE 求解器和根求解器）的梯度传播，实现对物理系统参数的优化。

提出的方法

• 实现根求解器、ODE 求解器（solve_ivp）和积分方法（mcquad, SQuad）等函数式，并采用解析导数以提升数值稳定性并减少内存消耗。
• 使用纯函数作为输入以确保高阶导数计算的正确性，通过跟踪影响输出的状态变量，将有状态对象方法转换为纯函数。
• 引入 EditableModule 基类，手动指定影响方法输出的参数名称，从而在复杂类中实现精确的梯度计算。
• 提供 LinearOperator 抽象以支持稀疏线性代数运算，避免在高维问题中进行昂贵的矩阵构造。
• 利用 PyTorch 的自动微分引擎结合解析导数，通过 gradgradcheck 验证，计算二阶及更高阶导数。
• 设计模块化库结构，包含优化、积分、插值和线性代数等专用模块，受 SciPy 启发，但专为可微分科学计算而定制。

实验结果

研究问题

• RQ1 是否能够通过解析导数实现可微分科学函数式，以提升科学模拟中的数值稳定性并减少内存使用？
• RQ2 如何高效且正确地计算复杂函数式（如 ODE 求解器和根求解器）的高阶导数（例如二阶导数）？
• RQ3 一个统一的可微分函数式库在多大程度上能够加速科学应用中的物理信息机器学习？
• RQ4 如何将有状态对象方法转换为纯函数，以在可微分科学计算中实现正确的梯度计算？
• RQ5 ξ-torch 是否能够通过端到端可微分模拟，有效实现物理系统（如镜面设计与分子动力学）的参数优化？

主要发现

• $\xi$-torch 库成功使用解析表达式计算科学函数式的高阶导数，相比反向传播至算法实现，数值稳定性更高且内存使用更低。
• 该库实现了正确的二阶梯度，经由 PyTorch 的 gradgradcheck 验证，证明其在复杂可微分函数式中具有鲁棒性。
• 在镜面设计案例中，根求解器函数式使基于神经网络的镜面表面在 200 次迭代内完成优化，成功将光束聚焦至目标点。
• 在分子动力学案例中，solve_ivp 函数式使初始速度的优化在 7500 次迭代内实现，使用 Adam 优化器成功形成字母 'OX'。
• 该库的模块化设计，以及对纯函数和可编辑状态变量的支持，使得在复杂物理模拟流水线中实现可靠的梯度计算成为可能。
• 可微分函数式（如蒙特卡洛积分与固定采样积分）的可用性，将可微分科学计算的适用范围扩展至标准 ODE 和根求解任务之外。