浅谈后向传递的计算量大约是前向传递的两倍

1. 前言

训练神经网络的一次迭代分为三步：（1）前向传递计算损失函数；（2）后向传递计算梯度；（3）优化器更新模型参数。在实验中，我们观察到一个现象： 后向传递的耗时几乎是前向传递的两倍，相比之下，优化器更新的耗时几乎可以忽略 。要解释这个现象，我们要从前向传递、后向传递和优化器参数更新的浮点数计算次数入手。

上图表示一次训练迭代中各个环节（前向传递、后向传递、通信环节、优化器更新）的耗时占比，来自于《PyTorch Distributed: Experiences on Accelerating Data Parallel Training》。上图中，纵轴表示耗时占比，FWD表示一次训练迭代中前向传递的耗时占比，BWD则表示一次训练迭代中后向传递的耗时占比，OPT表示一次训练迭代中优化器更新模型参数的耗时占比。从上图中可以看到，一次训练迭代中，后向传递的耗时几乎是前向传递的两倍，相比之下，优化器更新的耗时占比很小，几乎可以忽略。

上图表示GPipe流水线并行的调度策略，来自于《Efficient large-scale language model training on gpu clusters using megatron-lm》。上图中，横轴表示耗时，一个蓝色小块表示一个微批次的前向传递，一个绿色小块表示一个微批次的后向传递，黑色竖线表示一次流水线刷新，也就是优化器更新模型参数。从上图中可以看到，一个绿色小块的宽度大约是蓝色小块的二倍，一次训练迭代中，后向传递的耗时几乎是前向传递的两倍，相比之下，优化器更新的耗时占比很小，几乎可以忽略。

2. 反向传播算法是怎么工作的

反向传播算法已经是训练神经网络模型不可缺少的一部分。训练神经网络模型时，用梯度下降算法来学习和更新模型参数（包含权重weights和偏置bias），问题是如何计算损失函数关于模型参数的梯度呢？这就要用到反向传播（backpropagation）算法。

反向传播算法的核心是计算损失函数 关于神经网络权重 或偏置 的偏微分 或 ，这两个偏微分表达式就是梯度。梯度的物理意义是：我们改变网络权重和偏置可以在多大程度上影响损失函数 。

2.1 前向传递：计算神经网络的输出

在讨论反向传播算法之前，我们先以多层前馈神经网络为例，用基于矩阵的方法来计算神经网络的输出。在此过程中，先定义一些数学符号。

首先明确地定义神经网络中的权重参数，用 表示第 层第 个神经元到第 层第 个神经元的连接。例如，下图中的权重 表示第2层第4个神经元到第3层第2个神经元的连接。这个定义看起来有些麻烦，一个迷惑之处是 的顺序，应该用 表示输入神经元，用 表示输出神经元，而不是相反。但之后的计算中，你会看到这个定义是自然而然的。

类似地，我们定义网络中的偏置和激活。用 表示网络中第 层第 个神经元的偏置，用 表示网络中第 层第 个神经元的激活。下图给出了示例。

有了以上定义，我们就可以把网络第 层第 个神经元的激活 与第 层的激活联系起来。 单个神经元的计算公式(逐点形式):

其中，求和是对第 层上所有神经元 的求和， 是激活函数。

为了以矩阵形式重写上式，我们为第 层定义一个权重矩阵 (weight matrix) ，表示连接到第 层的权重，矩阵第 行第 列的元素为 。设第 层的输入维度为 ，输出维度为 。则该权重矩阵 的形状为 。

同样地，我们定义一个偏置向量 (bias vector) ，它的元素是 ，形状是 。定义一个激活向量 (activation vector) ，它的元素是 ，形状是 。

重写上式需要的最后一个要素是向量化 (vectorization) 函数。核心是函数逐点应用到向量中的每个元素。

有了以上定义，我们就可以 将上式改写为矩阵形式:

下标表示了矩阵或者向量的形状。这个公式给出一个更全局的视角来观察一层神经元的激活是如何与上一层神经元的激活联系起来的：先将权重矩阵应用到上一层的激活，再加上偏置向量，最后过激活函数。

我们 引入中间变量 ，我们将 称为第 层神经元的加权输入 (weighted input)。其元素为:

2.2 反向传播的四个基本等式

右边的第一项表示第 个输出激活多大程度上影响了损失函数。如果损失函数不太依赖于第 个神经元激活，那么 会是一个比较小的值。第二项表示激活函数在 上的改变程度。

将上式改写为矩阵形式:

证明:

从定义出发，应用多元微积分的链式法则:

2.2.2 误差与下一层误差的关联

上式中， 是第 层权重矩阵 的转置。这个式子乍一看是很复杂的，但每一项都有优美的解释。假设我们已经知道第 层的误差 ，乘以转置的权重矩阵 ，可以理解为将误差通过网络反向传播，这衡量了第 层输出激活的误差；再乘以阿达玛乘积 ，将误差通过第 层的激活函数反向传播，这衡量了第 层加权输入的误差。

这个式子是很优雅的。有了公式(BP1)和(BP2)，就可以计算出所有层对加权输入的误差了。 。这也就是称之为梯度反向传播算法的原因。

证明:

从定义出发，应用多元微积分的链式法则: 有 ， 以及 。

要评估等式右边的第一项，注意到：

取微分，得到:

带入上上式中，得到:

重写为矩阵形式，即:

2.2.3 偏置的梯度

第 层偏置的梯度就等于第 层加权输入的误差 。改写为矩阵形式:

证明:

由定义出发，用多元微积分的链式法则:

注意到: ，故 。带入上式中，得到:

2.2.4 权重的梯度

第 层权重的梯度就等于第 层加权输入的误差 与 上一层神经元输出激活 的乘积。将上式改写为:

假设权重 连接了两个神经元， 是神经元输入给权重的激活， 是权重输出给神经元的误差。

证明 ：

由定义出发，应用多元微积分的链式法则: