机器之心整理

参与：思、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* npdef 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 这样的通用自动微分库也许能在更广泛的领域发挥作用。