计算图上的微积分：反向传播

引言

反向传播是让深度模型训练在计算上变得可行的关键算法。对现代神经网络而言，相比朴素实现，它能让基于梯度下降的训练速度提升多达一千万倍。这就是“训练一周”和“训练二十万年”之间的差别。

反向传播不仅用于深度学习，它在很多其他领域里也是非常强大的计算工具，从天气预报到数值稳定性分析都能见到它，只是名字不同而已。事实上，这个算法至少在不同领域里被重复发明过几十次。更一般、与应用无关的名称是“反向模式自动微分”（reverse-mode differentiation）。

从根本上说，它是一种快速计算导数的技术。不只是做深度学习时需要它，在各种数值计算场景中，这也是非常值得掌握的技巧。

计算图

理解数学表达式时，计算图是一种很好的思考方式。比如考虑表达式 $e=(a+b)*(b+1)$。这里有三个操作：两次加法和一次乘法。为了方便讨论，我们引入两个中间变量 $c$ 和 $d$，让每个函数输出都对应一个变量：

构造计算图时，我们把这些操作以及输入变量都视为节点。若一个节点的值会作为另一个节点的输入，就从前者画一条箭头指向后者。

这类图在计算机科学中经常出现，尤其是在讨论函数式程序时。它和依赖图、调用图这些概念密切相关，也是很多深度学习框架背后的核心抽象。

我们可以给输入变量赋值，然后沿着图向上计算各个节点的值。例如设定 $a=2$、$b=1$：

这个表达式最终的值是 $6$。

计算图上的导数

如果想理解计算图中的导数，关键是先理解边上的导数。若 $a$ 会直接影响 $c$，我们就想知道它是如何影响 $c$ 的。也就是说，当 $a$ 发生一个微小变化时，$c$ 会怎么变化？这就是 $c$ 关于 $a$ 的偏导数。

为了计算这张图中的偏导数，我们需要用到求和法则和乘积法则：

下面这张图在每条边上都标出了对应导数。

如果我们想理解那些并不直接相连的节点之间是如何相互影响的，该怎么办？来看 $e$ 受 $a$ 影响的过程。若 $a$ 以速度 1 变化，那么 $c$ 也会以速度 $1$ 变化；接着，$c$ 以速度 $1$ 变化又会让 $e$ 以速度 $2$ 变化。因此，$e$ 相对于 $a$ 的变化率就是 $1*2$。

更一般地说，一个节点对另一个节点的影响，等于从前者到后者所有可能路径的贡献之和，而每条路径的贡献等于路径上各边导数的乘积。举例来说，若要求 $e$ 关于 $b$ 的导数，我们得到：

这同时考虑了 $b$ 通过 $c$ 影响 $e$，以及 $b$ 通过 $d$ 影响 $e$ 这两条路径。

这种“对所有路径求和”的一般规则，本质上只是多元链式法则的另一种视角。

路径分解

单纯“把所有路径都加起来”的问题在于，可能路径的数量非常容易出现组合爆炸。

在上图中，从 $X$ 到 $Y$ 有三条路径，从 $Y$ 到 $Z$ 也有三条路径。如果我们想通过对所有路径求和来计算 $\frac{\partial Z}{\partial X}$，就需要枚举 $3*3 = 9$ 条路径：

这里只是九条路径，但在更复杂的图里，路径数完全可能指数增长。

与其朴素地把所有路径都列出来，不如先做因式分解：

这正是“前向模式自动微分”和“反向模式自动微分”发挥作用的地方。它们本质上都是通过对路径做分解，来更高效地计算这类和式。与其显式枚举全部路径，不如在每个节点把路径重新合并。实际上，这两种算法都只会访问每条边一次。

前向模式自动微分从图的输入端出发，向末端推进。在每个节点，它把所有流入该节点的路径贡献累加起来。每条路径都代表输入影响该节点的一种方式，把它们加总之后，我们就得到输入对该节点的总影响，也就是它的导数。

即便你过去没有把它理解为图上的过程，如果学过微积分入门，前向模式自动微分其实和你隐含学过的很多求导思路非常接近。

与之相对，反向模式自动微分从图的输出端出发，向起点回传。在每个节点，它把所有从该节点发起的路径合并起来。

前向模式自动微分跟踪的是“某个输入如何影响每个节点”；反向模式自动微分跟踪的是“每个节点如何影响某个输出”。换句话说，前向模式是在每个节点上施加算子 $\frac{\partial}{\partial X}$，而反向模式是在每个节点上施加算子 $\frac{\partial Z}{\partial}$。¹

计算上的胜利

到这里，你可能会问：为什么有人会特别在意反向模式自动微分？它看起来只是以前向模式的另一种奇怪写法。它真的有优势吗？

我们再回到最初的那个例子：

如果我们从 $b$ 出发做前向模式自动微分，就能得到每个节点关于 $b$ 的导数。

这样我们算出了 $\frac{\partial e}{\partial b}$，也就是输出相对于某一个输入的导数。

那如果从 $e$ 向下做反向模式自动微分呢？这样我们会得到 $e$ 关于每个节点的导数：

这里说“关于每个节点”，字面意思就是每一个节点。我们不仅得到了 $\frac{\partial e}{\partial a}$，也得到了 $\frac{\partial e}{\partial b}$。前向模式只能给出“输出相对于单个输入”的导数；反向模式一次就能给出“输出相对于所有输入”的导数。

在这个小图里，这不过是两倍左右的提升。但设想一个函数有一百万个输入、一个输出。要用前向模式拿到全部输入的导数，你必须在图上跑一百万遍；而反向模式一次就能全部拿到。百万倍的速度提升，已经相当可观了。

训练神经网络时，我们把损失函数看作参数的函数。参数描述网络如何工作，而损失描述网络表现有多差。为了使用梯度下降，我们需要计算损失对所有参数的导数。现实中的神经网络常常拥有数百万，甚至上千万个参数。

所以，反向模式自动微分在神经网络语境下被称为“反向传播”，它带来的计算加速是压倒性的。

前向模式有没有更合适的时候？当然有。反向模式给的是“一个输出关于所有输入”的导数，而前向模式给的是“所有输出关于一个输入”的导数。如果函数的输出维度很大，前向模式有时会快得多。

这难道很平凡吗？

我第一次真正理解反向传播时，第一反应是：“这不就是链式法则吗？为什么人类花了这么久才想出来？”不只是我有这种感觉。表面上看，如果你问“前馈神经网络里有没有一种聪明的方式来计算导数”，答案似乎并不难。

但事情其实比它看上去难得多。在反向传播被发明出来的时代，人们关注的并不是我们今天熟悉的前馈神经网络；而且，当时也并不显然“用导数训练网络”是正确方向。只有当你先意识到导数能被快速算出来之后，这件事才显得理所当然。这里存在一种循环依赖。

更糟的是，只要粗略想一想，这个循环依赖中的任何一环都很容易被判定为“不可能”。用导数训练神经网络？那岂不是会卡在局部最小值里。计算这么多导数？那成本一定高得离谱。也正因为我们今天已经知道这条路是可行的，所以才不会本能地先列出一串它为何行不通的理由。

这就是事后视角的优势。一旦问题被正确框定，最难的工作往往已经完成了。

结语

这篇文章最重要的启示是：导数的计算成本比你想象得低。事实上，它低得有点反直觉，以至于人类在不同领域里一次又一次地重新发现了这件事。对深度学习来说，这是一个重要事实；对其他领域来说，只要这件事还不是常识，它就同样有价值。

除此之外，还有其他值得带走的东西。反向传播也是一副理解“导数如何在模型中流动”的透镜。这对于解释为什么某些模型难以优化非常有帮助。一个经典例子就是循环神经网络中的梯度消失问题。

最后，我认为这些技术背后还有更普遍的算法启示。反向传播和前向模式自动微分利用了一对非常强大的技巧：线性化与动态规划。借助它们，我们能以比直觉预期更高效的方式计算导数。如果你真正理解了这些技巧，还可以把它们用于很多其他涉及导数的有趣表达式。

这篇文章对反向传播的处理非常抽象。如果你想看一个更贴近神经网络本身的解释，仍然很值得继续读 Michael Nielsen 的相关章节。

致谢

感谢 Greg Corrado、Jon Shlens、Samy Bengio 和 Anelia Angelova 抽时间审阅原文。

也感谢 Dario Amodei、Michael Nielsen 和 Yoshua Bengio 参与关于“如何解释反向传播”的讨论。以及感谢所有容忍作者在讲座和研讨中反复练习这套解释的人。