自动微分技术在 AI for science 中的应用

本文简记我在学习自动微分相关技术时遇到的知识点。

以 NN 为代表的深度学习技术展现出了强大的参数拟合能力，人们通过堆叠固定的 layer 就能轻松设计出满足要求的参数拟合器。

例如，大部分图神经网络均基于消息传递的架构。在推理阶段，用户只需给出分子坐标及原子类型，就能得到整个分子的性质。因此其整体架构与下图类似：

在模型设计阶段，我们用 pytorch 即可满足大部分需求，以 schnetpack 为例：

我们 from torch import nn 导入了设计 nn 常用的模块。在初始化模型时，我们直接继承了 pytorch 内置的模块 class AtomisticModel(nn.Module)
有一些函数是重新编写的，例如激活函数 shiftedsoftplus

我们可以看到，模型的整体框架依然是基于 pytorch 的，但针对具体的应用场景，我们做了很多优化。

一方面，使用 pytorch 可以帮助我们快速建立类似上图的模型网络，pytorch 会自动执行梯度的反向传播。从 loss function 开始，逐层递进直至输入层。pytorch 还会帮助我们完成整个网络的参数迭代，学习率的迭代等等。。。

另一方面，针对一些特殊的需求，用户需要自行 DIY，完成需要的功能。

这其中隐含着，用户在程序设计时灵活性与便利性之间的折中。

注意到，刚才提到了梯度的反向传播，事实上，这种常用算法只是自动微分算法中的一种。引用 Gemini 的一个例子：

写到这里，自动微分技术的应用场景就很好理解了：

我们可以将参数优化的相关框架归结为两个应用场景：

针对第一个场景，我们可以使用 pytorch，因为 pytorch 对常用网络架构封装很好。

针对第二个场景，使用 pytorch 会更加繁琐，此时可以切换为 jax ，因为 jax 对用户自定义函数形式更加友好，其内置自动微分算法使用起来更加方便。

除了应用场景的区别外，二者还有以下几个区别：

余旷老师在他的系列博文里系统阐释了为什么 DMFF 要基于 jax 开发（参见：漫谈分子力场、自动微分与DMFF项目：4. DMFF和JAX概述）

总结一下，使用 jax 的原因有以下几点：

传统分子力场的形式不适合用 NN 建模
- 为方便大家理解，我举一个中学物理的例子。苹果从树上落下，遵从自由落体运动，位移随时间变化的规律：h=1/2 * g * t^2, 其中 g 作为引力常数就是需要通过多次落体实验测定的量。我们当然可以用多层 NN 拟合这一参数，但假如我们已经知道了这样一个表达式，此时直接使用该表达式即可。
- 传统分子力场就是高度参数化的方程，发展至今已经有了一套函数形式，无需从头用 NN 的形式拟合
反向传播算法只适用与 NN，不适应上述高度参数化的方程，但优化力场参数仍需要自动微分技术
- 计算原子受力，整个盒子的维里均需要微分技术，使用 jax 编程会更加方便
jax 性能更高，速度快
jax 可拓展性好
- 余旷老师在漫谈分子力场、自动微分与DMFF项目：5. DMFF中势函数的生成和拓展举了一个例子，使用 DMFF 能有效复用前人开发势函数模块，无需从头造轮子