BaiRuic
Go back

自动微分是什么

目录

周志华老师在《机器学习》中讲到， 机器学习算法包括 假设函数（模型）、损失函数（策略）、优化方法（算法） 这三部分， 深度学习模型自然也不例外。在模型训练时，计算损失值对假设函数中参数的梯度是最常见的操作之一。在PyTorch中，一句loss.backward()就计算好了参数的梯度了，一句optimizer.step()又可以更新参数，作为一个知其然也欲知其所以然的同学，我首先想知道他是怎么计算的微分？我将刨析其背后的细节。本文首先从微分的计算方式中引出了自动微分，然后又介绍了计算图这种工具，解释了如何基于计算图来实现自动微分。

众所周知，由于深度学习模型损失函数非凸，导致无法直接求得其解析解，因此很多模型中的参数是梯度下降不断迭代优化得到的，那在梯度下降的过程中，梯度是如何计算得到的呢？通常，梯度（即微分）的计算方式有如下几种：

数值微分（Numerical differentiation）

按照定义 直接计算 微分，如下：

∂f(θ)∂θi=lim⁡ϵ→0f(θ+ϵei)−f(θ)ϵ\frac{\partial f(\theta)}{\partial \theta_i}=\lim _{\epsilon \rightarrow 0} \frac{f\left(\theta+\epsilon e_i\right)-f(\theta)}{\epsilon}

其中 eie_i 为只有 ii 位置 为 11 其余均为 00 的向量。因为是计算 θ\theta 处对 θi\theta_i 的偏导，所以要保证除 θi\theta_i 之外的值不变。

但是一般为了更好的数值准确，会采用以下中心差分的方式计算：

∂f(θ)∂θi=f(θ+ϵei)−f(θ−ϵei)2ϵ+o(ϵ2)\frac{\partial f(\theta)}{\partial \theta_i}=\frac{f\left(\theta+\epsilon e_i\right)-f\left(\theta-\epsilon e_i\right)}{2 \epsilon}+o\left(\epsilon^2\right)

因为从泰勒展开可得

∂f(θ)∂θi=lim⁡ϵ→0f(θ+ϵei)−f(θ)ϵ+o(ϵ)\frac{\partial f(\theta)}{\partial \theta_i}=\lim _{\epsilon \rightarrow 0} \frac{f\left(\theta+\epsilon e_i\right)-f(\theta)}{\epsilon} + o\left(\epsilon \right)

相比之下，后者显然有更准确的梯度估计。

缺点：

  • 计算效率低下的缺陷。使用该方法，每计算一个参数的梯度，都需要前向传播两次，假设有nn个参数，则需要2n2n次前向传播。
  • 数值错误。体现在浮点运算的精度受限所导致的舍入误差（rounding error）；以及舍去高阶项导致的截断误差（truncation error）

因此，上述数值微分的方法更多用在梯度检验（大多深度学习框架均有使用），通常以如下形式展开检验：

δT∇θf(θ)=f(θ+ϵδ)−f(θ−ϵδ)2ϵ+o(ϵ2)\delta^T \nabla_\theta f(\theta)=\frac{f(\theta+\epsilon \delta)-f(\theta-\epsilon \delta)}{2 \epsilon}+o\left(\epsilon^2\right)

其中等号左边 ∇θf(θ)\nabla_\theta f(\theta) 为自动微分算法求解的梯度，δ\delta 为从单位球上随机选取的向量；等号右边为数值微分计算结果。通过判断等号左右两边数值来检验自动微分计算的梯度。

需要在注意的是：

  • δ\delta 需要尽可能覆盖所有梯度方向；
  • 等号右边 数值微分 是同时计算 θ\theta 中 所有维度的微分，因为微小变量 δ\delta 不是单一方向。

符号微分（Symbolic differentiation）

符号微分是除数值微分之外较为常见微分计算方式，在自动微分流行前，经常用该方法手工计算，因为可以保证计算精度。

用和、积、链式法则求梯度:

∂(f(θ)+g(θ))∂θ=∂f(θ)∂θ+∂g(θ)∂θ∂(f(θ)g(θ))∂θ=g(θ)∂f(θ)∂θ+f(θ)∂g(θ)∂θ∂f(g(θ))∂θ=∂f(g(θ))∂g(θ)∂g(θ)∂θ\begin{align*}& \frac{\partial(f(\theta)+g(\theta))}{\partial \theta}=\frac{\partial f(\theta)}{\partial \theta}+\frac{\partial g(\theta)}{\partial \theta}\\\\[1mm]& \frac{\partial(f(\theta) g(\theta))}{\partial \theta}=\mathrm{g}(\theta) \frac{\partial f(\theta)}{\partial \theta}+\mathrm{f}(\theta) \frac{\partial g(\theta)}{\partial \theta}\\\\[1mm]& \frac{\partial f(g(\theta))}{\partial \theta}=\frac{\partial f(g(\theta))}{\partial g(\theta)} \frac{\partial g(\theta)}{\partial \theta}\end{align*}

例如：

f(θ)=∏i=0nθif(θ)∂θk=∏j≠knθjf(\theta)=\prod_{i=0}^n \theta_i \quad \frac{f(\theta)}{\partial \theta_k}=\prod_{j \neq k}^n \theta_j

相较于数值微分，符号微分更为常见。然而，符号微分存在浪费计算资源的问题（即会有大量的重复计算和保存中间值的内存浪费）。例如上述示例中，计算 θ\theta 的微分，时间复杂度为 n(n-2) （n个维度，每个维度的偏导需要计算n-2个乘法）

自动微分 automatic differentiation

自动微分认为，任何数值计算的本质其实是一系列可微分算子的组合。那么，我们就可以假设我们求不出这个函数的导数，但是将该函数拆解成为其他子部分后，子部分可以通过常规的求导方式得到，最终将每个子部分进行组合，就得到了最终的结果。Therefore, 我们对一些常用的函数和表达式用像符号微分那样的求解方式求解，然后带入数值，作为中间结果进行保存，再将一系列保存下来的结果进行组合得到最终的目标值。由于在整个过程中，我们仅仅对一些基本函数或者特定表达式进行微分求解，那么就可以很容易和programming language中的for, if, while等结构进行组合，对用户完全隐藏了整个求解微分的细节。并且，他的计算本质上还是一种图计算，那么就可以对其进行一系列的优化来调优我们的系统。

在介绍自动微分之前，先介绍计算图这种工具。然后介绍 前向模式 和 反向模式 这两种自动微分机制。

计算图 Computational graph

Untitled 31

一张有向图，节点表示变量，边为变量之间的关系，即运算操作。计算的时候需要按照拓扑排序的顺序来计算。

前向模式的自动微分 Forward mode automatic differentiation (AD)

自动微分不是必须要反向传播，正向也可以计算微分。

前向自动微分计算过程

首先有 v1=1,v2=0\mathrm{v}_1=1, \mathrm{v}_2=0* ，*接着不断计算 vi\mathrm{v}{\mathrm{i}} 对 x1\mathrm{x}_1 的导数即可，中间可以通过链式法则进行计算，比如 v˙5=∂v5/∂v4∗v˙4\dot{v}_5= \partial \mathrm{v}_5 / \partial \mathrm{v}_4 * \dot{\mathrm{v}}_4 ，而之前已经求了 v˙2\dot{\mathrm{v}}_2 的结果，所以只需要计算当前一步的导数即可。

如果用数学语言来描述这个过程，就是需要计算 f\mathrm{f} 的 Jacobian 矩阵，其中 f:Rn→Rm\mathrm{f}: \mathcal{R}^{\mathrm{n}} \rightarrow \mathcal{R}^{\mathrm{m}} 表示由 n\mathrm{n} 个独立的输入变量 xi\mathrm{x}_{\mathrm{i}} 映射到 m\mathrm{m} 个相关的输出变量 yj\mathrm{y}_{\mathrm{j}} 。对于上面这种特殊的情况，可以把每一次 AutoDiff 的 foward pass 看成是将变量 x\mathrm{x} 的其中一个分量 x˙i=1\dot{x}_{\mathrm{i}}=1 其他的分量设为 00 的一次推导。所以当 f:R→Rm\mathrm{f}: \mathcal{R} \rightarrow \mathcal{R}^{\mathrm{m}} 时，forward pass 非常高效，因为所有 需要计算的偏导只需要进行一次 forward pass 即可。

反向模式的自动微分 Reverse mode automatic differentiation(AD)

反向模式的自动微分过程

左边的前向计算流程是一样的，但是在求导数的过程却是反向的，设定 vˉi=∂y∂vi\bar{v}_i=\frac{\partial y}{\partial v_i}，那么v7=1v_7 = 1，继续求v5,…v_5,…， 通过链式法则只需要计算当前一步的导数即可。

如果用数学语言来描述这个过程，就是需要计算 f\mathrm{f} 的 Jacobian 矩阵，其中 f:Rn→Rm\mathrm{f}: \mathcal{R}^{\mathrm{n}} \rightarrow \mathcal{R}^{\mathrm{m}} 。同样对于上面这种情况， 每一次 自动微分 的 backward pass 可以看成是将因变量 yy 的其中一个分量 yˉi=1\bar{y}_{\mathrm{i}}=1 其他分量设为 00 的一次推导。所以 当 f:Rn→Rf: \mathcal{R}^{\mathrm{n}} \rightarrow \mathcal{R} 时， reverse mode 非常高效，因为所有需要计算的偏导只需要进行一次 reverse pass 即可。 而我们知道在深度学习中 loss 一般都是一个标量，而参数一般都是一个高维张量，所以 f:Rn→R\mathrm{f}: \mathcal{R}^{\mathrm{n}} \rightarrow \mathcal{R} 可以表示绝大多 数深度学习模型的情况，通过上面的分析可以看出 reverse mode 效率更高，这也是为什么深度学习都是选择 reverse mode 进行梯度计算的原因，同时，这也是反向传播算法的由来。

注意事项：对于多路径的节点，例如下图 v1v_1 被多条路径使用:

Untitled 34

对 v1v_1 的微分就有：

v1‾=∂y∂v1=∂f(v2,v3)∂v2∂v2∂v1+∂f(v2,v3)∂v3∂v3∂v1=v2‾∂v2∂v1+v3‾∂v3∂v1\overline{v_1}=\frac{\partial y}{\partial v_1}=\frac{\partial f\left(v_2, v_3\right)}{\partial v_2} \frac{\partial v_2}{\partial v_1}+\frac{\partial f\left(v_2, v_3\right)}{\partial v_3} \frac{\partial v_3}{\partial v_1}=\overline{v_2} \frac{\partial v_2}{\partial v_1}+\overline{v_3} \frac{\partial v_3}{\partial v_1}

因此定义 vi→j‾=vj‾∂vj∂vi\overline{v_{i \rightarrow j}}=\overline{v_j} \frac{\partial v_j}{\partial v_i} 为输入输出节点对i→j 的伴随微分 （partial adjoint）。且有

vˉi=∑j∈next⁡(i)vi→j‾\bar{v}_i=\sum_{j \in \operatorname{next}(i)} \overline{v_{i \rightarrow j}}

即计算 viv_i 的微分，需要用到 节点 ii 所有的相邻伴随节点 jj 的微分。

注：partial adjoint 不知道应该翻译成什么，所以此处直接翻译为了伴随微分

反向自动微分算法理论

反向自动微分算法 Rverse AD algorithm

反向模式的自动微分伪代码。（注：上图中 return 语句应该与第一个for 对齐）

用一个字典来存放每一个节点的 partial adjoints。然后在计算每一个节点 ii 的时候，把所有后继节点 jj 的 partial adjoints** **相加即可。然后再计算节点 ii 的所有前驱节点 kk 的 partial adjoints。

通过扩展计算图实现反向模式的自动微分

Untitled 36

注：

  • 所谓的扩展计算图就是不在原先的计算图上表示，而是将 自动微分又表示出了一张计算图，这样实现的最大好处是容易计算 梯度的梯度，即二阶梯度，只需要对新的计算图再执行一次自动微分即可。
  • node_to_grad为二维列表，其中node_to_grad[i]存储了节点i的所有后继节点j的 vi→j‾\overline {v_{i→j}}  ，首先设定 out 的梯度为 1，out的前驱节点为节点4，因此有初始化node_to_grad[4]=[1]
  • 一个节点所有输出边梯度加在一起才是这个节点的梯度值，即 vˉi=∑j∈next⁡(i)vi→j‾\bar{v}_i=\sum_{j \in \operatorname{next}(i)} \overline{v_{i \rightarrow j}}

按照反拓扑序来遍历节点，计算到这个节点就代表着所有相关的梯度都计算完了。现在需要把相关的梯度都加起来，然后加起来的梯度作为这个节点的梯度。完整的计算过程如下：

  1. 首先节点 i = 4：
    1. 先计算节点 4 的微分，得到 v4‾\overline {v_4} = sum(node_to_grad[4]) = 1
    2. 然后遍历节点 4 的前驱节点 2 和 3，计算 v2→4‾\overline {v_{2→4}} 和 v3→4‾\overline {v_{3→4}}，并且分别存储到 node_to_grad[2] 和 node_to_grad[3] 中。具体结果为：v2→4‾=v3∗v4‾\overline {v_{2→4}} = v_3 * \overline{v_4}，v3→4‾=v4‾v2\overline {v_{3→4}} = \overline{v_4} v_2。
  2. 然后是节点 i = 3： 3. 先计算节点3的微分 v3‾\overline {v_3} = sum(node_to_grad[3]) = v3→4‾=v4‾v2\overline {v_{3→4}} = \overline{v_4} v_2 4. 然后遍历计算节点 3 的前驱节点 2 ，计算 v2→3‾\overline {v_{2→3}}，并把结果存储到 node_to_grad[2] 的列表中。具体计算得到 v2→3‾=v3‾∗1\overline {v_{2→3}} = \overline {v_3} * 1
  3. 然后是节点 i = 2： 5. 先计算 节点2的微分： v2‾\overline {v_2} = sum(node_to_grad[2]) = v2→3‾+v2→4‾\overline {v_{2→3}} + \overline {v_{2→4}} 6. 然后 遍历节点2的前驱节点 1，计算 v1→1‾\overline {v_{1→1}}，并把结果存储到 node_to_grad[1] 的列表中。具体计算得到 v1→2‾\overline {v_{1→2}} =v2‾∗exp(v1)= \overline {v_2} * exp(v_1)。
  4. 然后是节点 i = 1： 7. 先计算 节点1的微分： v1‾\overline {v_1} = sum(node_to_grad[1]) =v1→2‾\overline {v_{1→2}}。 8. 节点1没有前驱节点。

整个过程就是拓扑排序的逆向过程，即通过反拓扑序去生成一个反向的计算图。在反向自动微分的时候，会创建出一些中间新的节点，同时也会用到原先前向计算图中节点，最终得到完整的反向自动微分计算图，如上图右侧红色子图。

张量上的反向模式的自动微分

Untitled 37

与标量版本没有太大区别，只是需要注意 Tensor 的维度。

反向自动微分 vs 反向传播

Untitled 38

反向传播：

  • 先前向再回传，保留每个前向计算的中间结果用于计算梯度。
  • 只在一张计算图上执行，一方面我们要保留每个op的输入，这可能需要占相当大的memory；另一方面缺乏灵活性，例如我们需要计算梯度的梯度时就无法计算
  • 早期的神经网络框架如Caffe就是采用这种方式。

反向自动微分：

  • 反向自动微分的结果还是一张计算图。在新的计算图上可以有更多的操作，此外在新的计算图上再进行一次反向自动微分即可得到二阶梯度。
  • 先建立神经网络的计算图，再建立梯度的计算图。
  • 建立好计算图后再喂数据进行计算。只需要保留部分中间结果。现代神经网络框架如PyTorch采用的方式。

那么，自动微分具体是如何实现的呢？每个节点是如何保存张量和梯度信息的呢？将在自动微分是什么？(续) 中给出答案。

参考资料

  1. Automatic Differentiation in Machine Learning: a Survey
  2. What is Automatic Differentiation? - YouTube
Previous Post
自动微分是什么？(续)
Next Post
面向数据科学家的Docker指南