机器之心整理参与：思、Jamin一直以来，自动微分都在DL框架背后默默地运行着，本文希望探讨它到底是什么，通过JAX，自动微分又能怎么用。实战JAX自动微分Jax是谷歌开源的一个科学计算库，能对Python程序与NumPy运算执行自动微分，而且能够在GPU和TPU上运行，具有很高的性能。如下先导入JAX，然后用三行代码就能定义之前给出的反传不确定性度量。JAX这样的通用自动微分库也许能在更广泛的领域发挥作用。
机器之心整理
参与：思、Jamin
一直以来，自动微分都在 DL 框架背后默默地运行着，本文希望探讨它到底是什么，通过 JAX，自动微分又能怎么用。
自动微分现在已经是深度学习框架的标配，我们写的任何模型都需要靠自动微分机制分配模型损失信息，从而更新模型。在广阔的科学世界中，自动微分也是必不可少的。说到底，大多数算法都是由基本数学运算与基本函数组建的。
在 ICLR 2020 的一篇 Oral 论文中（满分 8/8/8），图宾根大学的研究者表示，目前深度学习框架中的自动微分模块只会计算批量数据反传梯度，但批量梯度的方差、海塞矩阵等其它量也很重要，它们可以在计算梯度的过程中快速算出来。
目前自动微分框架只计算出梯度，因此就限定了研究方向只能放在梯度下降变体之上，而不能做更广的探讨。为此，研究者构建了 BACKPACK，它建立在 PyTorch 之上，还扩展了自动微分与反向传播能获得的信息。
选自论文 BACKPACK，arXiv：1912.10985。
除此之外，Julia Computing 团队去年 7 月份也发表了一份论文，提出了可微编程系统，它能将自动微分内嵌于 Julia 语言，从而将其作为第一级的语言特性。由于广泛的科学计算和机器学习领域都需要线性代数的支持，因此这种可微编程能成为更加通用的一种模式。
从这些前沿研究可以清晰地感受到，自动微分越来越重要。
自动微分是什么
在数学与计算代数学中，自动微分也被称为微分算法或数值微分。它是一种数值计算的方式，用来计算因变量对某个自变量的导数。此外，它还是一种计算机程序，与我们手动计算微分的「分析法」不太一样。
自动微分基于一个事实，即每一个计算机程序，不论它有多么复杂，都是在执行加减乘除这一系列基本算数运算，以及指数、对数、三角函数这类初等函数运算。通过将链式求导法则应用到这些运算上，我们能以任意精度自动地计算导数，而且最多只比原始程序多一个常数级的运算。
一般而言会存在两种不同的自动微分模式，即前向累积梯度（前向模式）和反向累计梯度（反向模式）。前向累积会指定从内到外的链式法则遍历路径，即先计算 d_w1/d_x，再计算 d_w2/d_w1，最后计算 dy/dw_2。
反向梯度累积正好相反，它会先计算 dy/dw_2，然后计算 d_w2/d_w1，最后计算 d_w1/d_x。这是我们最为熟悉的反向传播模式，它非常符合「沿模型误差反向传播」这一直观思路。
如图所示，两种自动微分模式都在求 dy/dx，只不过根据链式法则展开的形式不太一样。
来一个实例：误差传播
在统计学上，由于变量含有误差，使得函数也含有误差，我们将其称之为误差传播。阐述这种关系的定律叫做误差传播定律。
先定义一个函数 q(x,y) ，我们想通过 q 传递 x 与 y 的不确定性信息，即 _x 与 _y。最直接的方式是随机采样 x 与 y，并计算 q 的值，然后查看它的分布。这就是「传播不确定性」这个概念的意义。
误差传播的积分公式可以是一个近似值， q(x,y) 的一般表达式可以写为：
如果我们定义一个特殊案例，即 q(x,y)=x±y，那么总不确定性可以写为：
对于特例 q(x,y)=xy 与 q(x,y)=x/y ，不确定性分别为 (σ_q/q)^2 = (σ_x/x)^2 (σ_y/y)^2 与 σ_q=(x/y)* sqrt((σ_x/x)^2 (σ_y/y)^2)。
我们可以尝试这些方法，并对比根据这些近似公式算出来的反传误差，以及实际发生的反传误差。
实战 JAX 自动微分
Jax 是谷歌开源的一个科学计算库，能对 Python 程序与 NumPy 运算执行自动微分，而且能够在 GPU 和 TPU 上运行，具有很高的性能。
如下先导入 JAX，然后用三行代码就能定义之前给出的反传不确定性度量。
from jax *import* grad, jacfwd
import jax.numpy *as* np
def error_prop_jax_gen(q,x,dx):
jac = jacfwd(q)
return np.sqrt(np.sum(np.power(jac(x)*dx,2)))
这里计算的反传梯度是根据 jax 完成的，后面的反传误差会直接通过公式计算，并对比两者。
1. 配置两个具有不确定性的观察值
我们需要使用 x 与 y 作为符号推理，但可以把它们都储存在数组 x 中，x[0]=x、x[1]=y。
x_ = np.array([2.,3.])
dx_ = np.array([.1,.1])
2. 加减法
在 (,)=± 这一特例情况下，误差传播公式可以简化为
上图所示，通过误差传播公式计算出来的值与 JAX 计算出来的是一致地。
3. 乘除法
在 (,)= 与 (,)=/ 这两种特例中，误差传播公式可以写为：
4. 幂
对于特例 (,)=^*^，传播公式可以表示为：
我们可以写成
JAX 的使用非常多样，甚至能直接使用它搭建神经网络。例如 JAXnet 框架，它是一个基于 JAX 的深度学习库，它的 API 提供了便利的模型搭建体验。比如说，以下代码就能建个神经网络：
from jaxnet import *
net = Sequential(Dense(1024), relu, Dense(1024), relu, Dense(4), logsoftmax)
此外，不久之前，DeepMind 也发布了两个新库：在 Jax 上进行面向对象开发 的 Haiku 和 Jax 上的强化学习库 RLax。JAX 这样的通用自动微分库也许能在更广泛的领域发挥作用。
