MindSpore自动微分快速教程

MindSpore拥有完善的自动微分系统。本文将会借着对自动微分思想的介绍来展示MindSpore自动微分的各项能力，方便读者运用在自己的项目中。

grad 的入参为一个函数，返回的是求导后的函数。定义一个Python函数f用来计算数学函数$f$，grad(f) 就是表达$\nabla f$的Python函数。 grad(f)(x) 就是$\nabla f(x)$的值。

由于 grad 作用在函数上，所以grad也可以用来处理它自己的输出：

下面是一个计算线性回归模型的梯度的例子，首先：

在grad 中使用 grad_position对指定的位置参数进行微分

本质上来说，使用grad_position时，如果f是一个Python函数，那么表达式grad(f, i)就是在求偏微分$\partial_i f$.

value_and_grad:同时获得函数值与梯度

value_and_grad可以方便地同时计算函数值和梯度值：

与数值计算结果比较

自动微分可以很直接地与有限微分比较：

grad+grad得到Hessian向量积

使用高阶grad可以构造Hessian向量积。(后面我们会用前向模式和反向模式写一个更高效的实现)

Hessian向量积可用来在截断的牛顿共轭梯度法中最小化一个光滑的凸函数，或者用来判断神经网络训练目标的曲率性质。(如 1, 2, 3, 4).

对于一个有着连续二阶导的标量函数（这种函数的Hessian矩阵是对称的）$f : \mathbb{R}^n \to \mathbb{R}$，点$x \in \mathbb{R}^n$处的Hessian算子为$\partial^2 f(x)$。一个Hessian向量积用来计算映射：

$\qquad v \mapsto \partial^2 f(x) \cdot v$

其中 $v \in \mathbb{R}^n$。

有一个技巧是我们不能实例化整个Hessian矩阵：如果$n$很大的话（神经网络中可能达到百万或上亿的量级），完整的Hessian矩阵是没法存储的。

幸运的是, grad 提供了一种高效计算Hessian向量积的方式。我们只需要有恒等式：

$\qquad \partial^2 f (x) v = \partial [x \mapsto \partial f(x) \cdot v] = \partial g(x)$，

其中 $g(x) = \partial f(x) \cdot v$ 是一个新的标量函数，其表示 $f$在$x$的梯度与向量$v$的点乘。这里只涉及对标量函数的向量值的微分，这种情形下 grad 是高效的。

用MindSpore代码，我们可以写出：

这个例子表明我们可以自由的使用词汇闭包，MindSpore都可以正确处理。在后面我会看到Hessian矩阵是怎么被计算出来的，知晓了原理之后我们会同时运用前向模式和反向模式提供一个更高效的写法。

运用 jacfwd 和 jacrev 计算Jacobians 和 Hessians 矩阵

用户可以用 jacfwd 和 jacrev计算Jacobian矩阵：

这两个函数得到的结果应该是一样的，二者只是实现方式不通： jacfwd 使用的是前向模式的自动微分，在比较"高"的Jacobian矩阵上较高效。 jacrev 使用的是反向模式，在"宽"的矩阵上更高效。对于比较方正的矩阵， jacfwd 效果稍好于jacrev。

关于前向模式和反向模式的更多信息，请继续阅读！

使用一种组合的方式计算dense的Hessian矩阵：

这里的shape是合理的：$f : \mathbb{R}^n \to \mathbb{R}^m$, 在点 $x \in \mathbb{R}^n$ 上，会有shape

• $f(x) \in \mathbb{R}^m$, $f$ 在 $x$ 处的值，
• $\partial f(x) \in \mathbb{R}^{m \times n}$, $x$ 处的Jacobian矩阵，
• $\partial^2 f(x) \in \mathbb{R}^{m \times n \times n}$, $x$ 处的Hessian矩阵

jacfwd(jacrev(f)) 或 jacrev(jacfwd(f)) 或者二者任意的组合皆可实现一个hessian矩阵，只是 forward+reverse一般情况下是效率最高的方式。 这是因为里面一层的Jacobian计算经常会有针对宽Jacobian矩阵的微分（比如loss function $f : \mathbb{R}^n \to \mathbb{R}$），在外面那一层的Jacobian 计算 通常是微分一个正方矩阵（因为会有$\nabla f : \mathbb{R}^n \to \mathbb{R}^n$），这时forward-mode速度更快。

深入理解两个基本的自动微分函数

Jacobian向量积 (JVPs, 前向模式自动微分)

MindSpore对前向和反向的自动微分都提供了高效且泛用性强的实现。我们熟悉的 grad 是基于反向模式实现的，不过为了理解二者的区别，我们需要一点数学背景。

JVPs的数学背景

从数学的角度看，给定一个函数 $f : \mathbb{R}^n \to \mathbb{R}^m$，$f$ 在输入点 $x \in \mathbb{R}^n$ 的Jacobian矩阵可被记作 $\partial f(x)$，通常型如 $\mathbb{R}^m \times \mathbb{R}^n$:

$\qquad \partial f(x) \in \mathbb{R}^{m \times n}$.

我们可以将 $\partial f(x)$ 视为线性映射，把在点 $x$ 处 $f$ 定义域上的正切空间（ 其实就是 $\mathbb{R}^n$ 的一份拷贝）映射到了在点 $f(x)$ 处 $f$ 陪域上的正切空间（$\mathbb{R}^m$ 的拷贝）。

$\qquad \partial f(x) : \mathbb{R}^n \to \mathbb{R}^m$.

这个映射又被称作 $f$ 在 $x$ 的前推映射。Jacobian矩阵只是这个线性映射在标准情况下的矩阵形式。

如果我们不拘泥于一个特定的点 $x$，那么函数 $\partial f$ 可被视为先取一个输入点然后返回那个点上的Jacobian线性映射：

$\qquad \partial f : \mathbb{R}^n \to \mathbb{R}^n \to \mathbb{R}^m$.

尤其是，做反curring时，给定输入 $x \in \mathbb{R}^n$ 和切向量 $v \in \mathbb{R}^n$，返回一个输出切向量 $\mathbb{R}^m$。我们把从 $(x, v)$ 到输出切向量的映射称之为 Jacobian向量积，写作：

$\qquad (x, v) \mapsto \partial f(x) v$

MindSpore中的JVP

回到Python代码上，MindSpore的 jvp 函数模拟了上述转换。 给定一个Python函数 $f$, MindSpore的 jvp 可以得到一个表达 $(x, v) \mapsto (f(x), \partial f(x) v)$ 的函数

按照Haskell类型风格, 有：

jvp :: (a -> b) -> a -> T a -> (b, T b)

在这里，我们用 T a 表示 a 切空间的类型。简而言之， jvp 的参数有 a -> b类型函数,、 a 类型的值和T a切向量。返回的是b类型的值和T b 类型的切向量。

jvp的计算方式与原始函数很相似，但它与每个a类型的原始值配对，并推送T a类型的切线值。对于每个原始函数会应用的原始数值操作，jvp转换后的函数为该原始函数执行一个 "JVP规则"，既对原始值进行评估，又在这些原始值上应用原始的JVP。

这种计算策略对计算的复杂度有直接的影响：因为在计算JVP的过程中不用存储任何东西，所以空间开销和计算的深度完全无关。除此之外， jvp 转换过的函数FLOP开销约是原函数的3倍 （一份来自原函数的计算，比如 sin(x)； 一份来自线性化，如 cos(x)；还有一份来自于将线性化函数施加在向量上，如 cos_x * v）。 换句话说，对于固定的点 $x$，我们计算 $v \mapsto \partial f(x) \cdot v$ 和计算 $f$的边际成本是相近的。

这里的空间复杂度看起来很有说服力，但我们在机器学习中并不经常见到前向模式。

为了回答这个问题，首先假设要用JVP构建一个完整的Jacobian矩阵。如果我们是对一个one-hot切向量用了JVP，结果反映的是Jacobian矩阵的一列，对应填入的非零项。所以我们是可以通过一次构建一列的方式构建一个完整的Jacobian矩阵的，而且每一列的开销和一次函数计算差不多。这就意味这对于"高"的Jaocbian矩阵来说比较合算，但对于"宽"的就较为低效。

如果在机器学习中做基于梯度的优化，你可能想要最小化损失函数，这个损失函数以 $\mathbb{R}^n$ 为参数，返回一个标量值$\mathbb{R}$。 这就意味着该函数的Jacobian矩阵会很宽了：$\partial f(x) \in \mathbb{R}^{1 \times n}$，一般我们会认为和梯度向量 $\nabla f(x) \in \mathbb{R}^n$ 一样。一次一列地构建这个矩阵，而且每列的FLOP和原函数计算一次的开销差不多，这个开销当然是不小的。尤其是，对于训练神经网络来说，损失函数 $f$ 的 $n$ 可以达到上亿的量级，这就更暴露出前向模式的问题了。

为了解决这种问题，就需要反向模式了。

向量Jacobian 积 (VJP, 反向模式自动微分)

和前向模式的一次一列的方式不同，反向模式的构造方式是一次一行。

VJPs 的数学背景

首先考虑有 $f : \mathbb{R}^n \to \mathbb{R}^m$。 其VJP表达为：

$\qquad (x, v) \mapsto v \partial f(x)$,

其中 $v$ 是 $f$ 在 $x$ 的余切空间($\mathbb{R}^m$ 的同构)。严谨来说，$v$ 是线性映射 $v : \mathbb{R}^m \to \mathbb{R}$， $v \partial f(x)$ 指的是复合函数 $v \circ \partial f(x)$，在 $\partial f(x) : \mathbb{R}^n \to \mathbb{R}^m$ 时成立。 不过通常 $v$ 都可以视为 $\mathbb{R}^m$ 中的向量，这两个写法基本可以互换。

有了这些说明后，我们把VJP的线性部分视为JVP线性部分的转置（或伴随、共轭）：

$\qquad (x, v) \mapsto \partial f(x)^\mathsf{T} v$.

对点 $x$，有：

$\qquad \partial f(x)^\mathsf{T} : \mathbb{R}^m \to \mathbb{R}^n$.

对余切空间的映射通常称为 $f$ 在 $x$ 的拉回。理解的关键在于拉回会从形似输出 $f$ 的形式得到形似输入 $f$ 的形式，就像线性函数转置一样。

MindSpore中使用VJP

MindSpore vjp 以一个python函数 $f$ 为输入，返回表示 VJP $(x, v) \mapsto (f(x), v^\mathsf{T} \partial f(x))$ 的函数。

仿照 Haskell-like type signatures, 有

vjp :: (a -> b) -> a -> (b, CT b -> CT a)

其中，我们用CT a来表示a的余切空间的类型。换句话说，vjp将一个a -> b类型的函数和一个a类型的点作为参数，并返回一个由b类型的值和CT b -> CT a类型的线性映射组成的对。

VJP一个优良的性质在于VJP是按行构建Jacobian矩阵， $(x, v) \mapsto (f(x), v^\mathsf{T} \partial f(x))$ 的FLOP仅为计算 $f$ 的三倍左右。而且计算 $f : \mathbb{R}^n \to \mathbb{R}$ 的梯度，我们只需要一次VJP就够了。这就是为什么 grad 在大的神经网络中做梯度优化依然高效。

不过还有一点需要考虑一下： 尽管 FLOP开销不高，VJP的空间复杂度是随计算深度上升而上升的。而且实现上通常比前向模式复杂。

反向模式的更多说明请参阅 this tutorial video from the Deep Learning Summer School in 2017.

VJP计算梯度向量

可以用VJP得到梯度向量：

用前向和反向模式得到Hessian向量积

仅用反向模式得到Hessian向量积的实现：

通过组合使用前反向的方法我们可以得到更高效的实现。

设有待微分函数 $f : \mathbb{R}^n \to \mathbb{R}$ ， 在点 $x \in \mathbb{R}^n$ 线性化函数，并有向量 $v \in \mathbb{R}^n$。 Hessian向量积函数为：

$(x, v) \mapsto \partial^2 f(x) v$

构造helper function $g : \mathbb{R}^n \to \mathbb{R}^n$，定义为 $f$ 的导数（或梯度）, 即 $g(x) = \partial f(x)$。使用一次JVP，便得到：

$(x, v) \mapsto \partial g(x) v = \partial^2 f(x) v$。

用代码写作：

在这里我们不需要 ops.inner，该 hvp 函数对任何shape的数组都成立。

以下是该函数的一个样例：

你也可以考虑写一种先前向后反向的方式：

不过这就不是很高效了，因为前向模式比反向模式的开销低一些，而且由于外层微分算子计算量比内层的要大，继续在外层用前向模式反而更好：