周志华老师在《机器学习》中讲到, 机器学习算法包括 假设函数(模型)、损失函数(策略)、优化方法(算法) 这三部分, 深度学习模型自然也不例外。在模型训练时,计算损失值对假设函数中参数的梯度是最常见的操作之一。在PyTorch中,一句loss.backward()就计算好了参数的梯度了,一句optimizer.step()又可以更新参数,作为一个知其然也欲知其所以然的同学,我首先想知道他是怎么计算的微分?我将刨析其背后的细节。本文首先从微分的计算方式中引出了自动微分,然后又介绍了计算图这种工具,解释了如何基于计算图来实现自动微分。
众所周知,由于深度学习模型损失函数非凸,导致无法直接求得其解析解,因此很多模型中的参数是梯度下降不断迭代优化得到的,那在梯度下降的过程中,梯度是如何计算得到的呢?通常,梯度(即微分)的计算方式有如下几种:
数值微分(Numerical differentiation)
按照定义 直接计算 微分,如下:
$$ \frac{\partial f(\theta)}{\partial \theta_i}=\lim _{\epsilon \rightarrow 0} \frac{f\left(\theta+\epsilon e_i\right)-f(\theta)}{\epsilon} $$
其中 $e_i$ 为只有 $i$ 位置 为 $1$ 其余均为 $0$ 的向量。因为是计算 $\theta$ 处对 $\theta_i$ 的偏导,所以要保证除 $\theta_i$ 之外的值不变。
但是一般为了更好的数值准确,会采用以下中心差分的方式计算:
$$ \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) $$
因为从泰勒展开可得
$$ \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) $$
相比之下,后者显然有更准确的梯度估计。
缺点:
- 计算效率低下的缺陷。使用该方法,每计算一个参数的梯度,都需要前向传播两次,假设有$n$个参数,则需要$2n$次前向传播。
- 数值错误。体现在浮点运算的精度受限所导致的舍入误差(rounding error);以及舍去高阶项导致的截断误差(truncation error)
因此,上述数值微分的方法更多用在梯度检验(大多深度学习框架均有使用),通常以如下形式展开检验:
$$ \delta^T \nabla_\theta f(\theta)=\frac{f(\theta+\epsilon \delta)-f(\theta-\epsilon \delta)}{2 \epsilon}+o\left(\epsilon^2\right) $$
其中等号左边 $\nabla_\theta f(\theta)$ 为自动微分算法求解的梯度,$\delta$ 为从单位球上随机选取的向量;等号右边为数值微分计算结果。通过判断等号左右两边数值来检验自动微分计算的梯度。
需要在注意的是:
- $\delta$ 需要尽可能覆盖所有梯度方向;
- 等号右边 数值微分 是同时计算 $\theta$ 中 所有维度的微分,因为微小变量 $\delta$ 不是单一方向。
符号微分(Symbolic differentiation)
符号微分是除数值微分之外较为常见微分计算方式,在自动微分流行前,经常用该方法手工计算,因为可以保证计算精度。
用和、积、链式法则求梯度:
$$ \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(\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
一张有向图,节点表示变量,边为变量之间的关系,即运算操作。计算的时候需要按照拓扑排序的顺序来计算。
前向模式的自动微分 Forward mode automatic differentiation (AD)
自动微分不是必须要反向传播,正向也可以计算微分。
首先有 $\mathrm{v}_1=1, \mathrm{v}_2=0$* ,*接着不断计算 $\mathrm{v}{\mathrm{i}}$ 对 $\mathrm{x}_1$ 的导数即可,中间可以通过链式法则进行计算,比如 $\dot{v}_5= \partial \mathrm{v}_5 / \partial \mathrm{v}_4 * \dot{\mathrm{v}}_4$ ,而之前已经求了 $\dot{\mathrm{v}}_2$ 的结果,所以只需要计算当前一步的导数即可。
如果用数学语言来描述这个过程,就是需要计算 $\mathrm{f}$ 的 Jacobian 矩阵,其中 $\mathrm{f}: \mathcal{R}^{\mathrm{n}} \rightarrow \mathcal{R}^{\mathrm{m}}$ 表示由 $\mathrm{n}$ 个独立的输入变量 $\mathrm{x}{\mathrm{i}}$ 映射到 $\mathrm{m}$ 个相关的输出变量 $\mathrm{y}{\mathrm{j}}$ 。对于上面这种特殊的情况,可以把每一次 AutoDiff 的 foward pass 看成是将变量 $\mathrm{x}$ 的其中一个分量 $\dot{x}_{\mathrm{i}}=1$ 其他的分量设为 $0$ 的一次推导。所以当 $\mathrm{f}: \mathcal{R} \rightarrow \mathcal{R}^{\mathrm{m}}$ 时,forward pass 非常高效,因为所有 需要计算的偏导只需要进行一次 forward pass 即可。
反向模式的自动微分 Reverse mode automatic differentiation(AD)
左边的前向计算流程是一样的,但是在求导数的过程却是反向的,设定 $\bar{v}_i=\frac{\partial y}{\partial v_i}$,那么$v_7 = 1$,继续求$v_5,…$, 通过链式法则只需要计算当前一步的导数即可。
如果用数学语言来描述这个过程,就是需要计算 $\mathrm{f}$ 的 Jacobian 矩阵,其中 $\mathrm{f}: \mathcal{R}^{\mathrm{n}} \rightarrow \mathcal{R}^{\mathrm{m}}$ 。同样对于上面这种情况, 每一次 自动微分 的 backward pass 可以看成是将因变量 $y$ 的其中一个分量 $\bar{y}_{\mathrm{i}}=1$ 其他分量设为 $0$ 的一次推导。所以 当 $f: \mathcal{R}^{\mathrm{n}} \rightarrow \mathcal{R}$ 时, reverse mode 非常高效,因为所有需要计算的偏导只需要进行一次 reverse pass 即可。 而我们知道在深度学习中 loss 一般都是一个标量,而参数一般都是一个高维张量,所以 $\mathrm{f}: \mathcal{R}^{\mathrm{n}} \rightarrow \mathcal{R}$ 可以表示绝大多 数深度学习模型的情况,通过上面的分析可以看出 reverse mode 效率更高,这也是为什么深度学习都是选择 reverse mode 进行梯度计算的原因,同时,这也是反向传播算法的由来。
注意事项:对于多路径的节点,例如下图 $v_1$ 被多条路径使用:
对 $v_1$ 的微分就有:
$$ \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} $$
因此定义 $\overline{v_{i \rightarrow j}}=\overline{v_j} \frac{\partial v_j}{\partial v_i}$ 为输入输出节点对i→j 的伴随微分 (partial adjoint)。且有
$$ \bar{v}i=\sum{j \in \operatorname{next}(i)} \overline{v_{i \rightarrow j}} $$
即计算 $v_i$ 的微分,需要用到 节点 $i$ 所有的相邻伴随节点 $j$ 的微分。
注:partial adjoint 不知道应该翻译成什么,所以此处直接翻译为了伴随微分
反向自动微分算法理论
反向自动微分算法 Rverse AD algorithm
用一个字典来存放每一个节点的 partial adjoints。然后在计算每一个节点 $i$ 的时候,把所有后继节点 $j$ 的 partial adjoints** **相加即可。然后再计算节点 $i$ 的所有前驱节点 $k$ 的 partial adjoints。
通过扩展计算图实现反向模式的自动微分
注:
- 所谓的扩展计算图就是不在原先的计算图上表示,而是将 自动微分又表示出了一张计算图,这样实现的最大好处是容易计算 梯度的梯度,即二阶梯度,只需要对新的计算图再执行一次自动微分即可。
node_to_grad为二维列表,其中node_to_grad[i]存储了节点i的所有后继节点j的 $\overline {v_{i→j}}$ ,首先设定 out 的梯度为 1,out的前驱节点为节点4,因此有初始化node_to_grad[4]=[1]- 一个节点所有输出边梯度加在一起才是这个节点的梯度值,即 $\bar{v}i=\sum{j \in \operatorname{next}(i)} \overline{v_{i \rightarrow j}}$
按照反拓扑序来遍历节点,计算到这个节点就代表着所有相关的梯度都计算完了。现在需要把相关的梯度都加起来,然后加起来的梯度作为这个节点的梯度。完整的计算过程如下:
- 首先节点
i = 4:- 先计算节点 4 的微分,得到 $\overline {v_4}$ =
sum(node_to_grad[4]) = 1 - 然后遍历节点 4 的前驱节点 2 和 3,计算 $\overline {v_{2→4}}$ 和 $\overline {v_{3→4}}$,并且分别存储到
node_to_grad[2]和node_to_grad[3]中。具体结果为:$\overline {v_{2→4}} = v_3 * \overline{v_4}$,$\overline {v_{3→4}} = \overline{v_4} v_2$。
- 先计算节点 4 的微分,得到 $\overline {v_4}$ =
- 然后是节点
i = 3: 3. 先计算节点3的微分 $\overline {v_3}$= sum(node_to_grad[3]) =$\overline {v_{3→4}} = \overline{v_4} v_2$ 4. 然后遍历计算节点 3 的前驱节点 2 ,计算 $\overline {v_{2→3}}$,并把结果存储到node_to_grad[2]的列表中。具体计算得到 $\overline {v_{2→3}} = \overline {v_3} * 1$ - 然后是节点
i = 2: 5. 先计算 节点2的微分: $\overline {v_2}$ =sum(node_to_grad[2])= $\overline {v_{2→3}} + \overline {v_{2→4}}$ 6. 然后 遍历节点2的前驱节点 1,计算 $\overline {v_{1→1}}$,并把结果存储到node_to_grad[1]的列表中。具体计算得到 $\overline {v_{1→2}}$ $= \overline {v_2} * exp(v_1)$。 - 然后是节点
i = 1: 7. 先计算 节点1的微分: $\overline {v_1}$ =sum(node_to_grad[1])=$\overline {v_{1→2}}$。 8. 节点1没有前驱节点。
整个过程就是拓扑排序的逆向过程,即通过反拓扑序去生成一个反向的计算图。在反向自动微分的时候,会创建出一些中间新的节点,同时也会用到原先前向计算图中节点,最终得到完整的反向自动微分计算图,如上图右侧红色子图。
张量上的反向模式的自动微分
与标量版本没有太大区别,只是需要注意 Tensor 的维度。
反向自动微分 vs 反向传播
反向传播:
- 先前向再回传,保留每个前向计算的中间结果用于计算梯度。
- 只在一张计算图上执行,一方面我们要保留每个op的输入,这可能需要占相当大的memory;另一方面缺乏灵活性,例如我们需要计算梯度的梯度时就无法计算
- 早期的神经网络框架如Caffe就是采用这种方式。
反向自动微分:
- 反向自动微分的结果还是一张计算图。在新的计算图上可以有更多的操作,此外在新的计算图上再进行一次反向自动微分即可得到二阶梯度。
- 先建立神经网络的计算图,再建立梯度的计算图。
- 建立好计算图后再喂数据进行计算。只需要保留部分中间结果。现代神经网络框架如PyTorch采用的方式。
那么,自动微分具体是如何实现的呢?每个节点是如何保存张量和梯度信息的呢?将在自动微分是什么?(续) 中给出答案。