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

Author: [回旋托马斯x]

Link: [https://zhuanlan.zhihu.com/p/675517271]

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）算法。

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

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

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

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

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

有了以上定义，我们就可以把网络第$l$层第$j$个神经元的激活 $a_j^l$ 与第 $l-1$ 层的激活联系起来。单个神经元的计算公式(逐点形式)： $a_j^l=\sigma (\sum _k{w}_{jk}^l{a}_k^{l-1}+b_j^l)$

其中，求和是对第$l-1$层上所有神经元 $k$ 的求和， $\sigma ()$ 是激活函数。

为了以矩阵形式重写上式，我们为第$l$层定义一个权重矩阵（weight matrix） $w^l$ ，表示连接到第$l$层的权重，矩阵第$j$行第$k$列的元素为 $w_{jk}^l$ 。设第$l$层的输入维度为 $M$ ，输出维度为 $N$ 。则该权重矩阵 $w^l$ 的形状为 $N\times M$ 。

同样地，我们定义一个偏置向量（bias vector） $b^l$ ，它的元素是 $b_j^l$ ，形状是 $N$ 。定义一个激活向量（activation vector） $a^l$ ，它的元素是 $a_j^l$ ，形状是 $N$ 。

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

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

$a_N^l=\sigma (w_{N\times M}^l{a}_M^{l-1}+b_N^l)$

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

我们引入中间变量 $z^l=w^l{a}^{l-1}+b^l$ ，我们将 $z^l$ 称为第 $l$ 层神经元的加权输入（weighted input）。其元素为：

$z_{jk}^l=\sum _k{w}_{jk}^l{a}_k^{l-1}+b_j^l$

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

反向传播是要理解改变网络中的权重和偏置多大程度上影响损失函数。这意味着计算偏微分 $\frac{\partial C}{\partial w_{jk}^l}$ 和 $\frac{\partial C}{\partial b_j^l}$ 。为了计算这两个偏微分，我们引入中间变量 ${\delta }_j^l$ ，称为第$l$层第$j$个神经元的误差（error）。反向传播会给出计算误差 ${\delta }_j^l$ 的步骤，再将 ${\delta }_j^l$ 与 $\frac{\partial C}{\partial w_{jk}^l}$ 和 $\frac{\partial C}{\partial b_j^l}$ 联系起来。

将误差 ${\delta }_j^l$ 定义为： ${\delta }_j^l=\frac{\partial C}{\partial z_j^l}$

2.2.1 输出层的误差

${\delta }_j^L=\frac{\partial C}{\partial a_j^L}{\sigma }^{{}^{\prime }}(z_j^L)\cdots \cdots (BP1)$ 右边的第一项表示第$j$个输出激活多大程度上影响了损失函数。如果损失函数不太依赖于第$j$个神经元激活，那么 ${\delta }_j^L$ 会是一个比较小的值。第二项表示激活函数在 $z_j^L$ 上的改变程度。

将上式改写为矩阵形式： ${\delta }^L={\mathrm{\nabla }}_{a}C\odot {\sigma }^{{}^{\prime }}(z^L)$

证明：

从定义出发，应用多元微积分的链式法则： ${\delta }_j^L=\frac{\partial C}{\partial z_j^L}=\sum _k\frac{\partial C}{\partial a_k^L}\frac{\partial a_k^L}{z_j^L}=\frac{\partial C}{\partial a_j^L}\frac{\partial a_j^L}{z_j^L}=\frac{\partial C}{\partial a_j^L}{\sigma }^{{}^{\prime }}(z_j^L)$

2.2.2 误差与下一层误差的关联\delta^{l}=((w^{l+1})^{T}\delta^{l+1})\odot\sigma^{'}(z^l) \qquad\qquad\qquad \cdots\cdots(BP2) \

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

这个式子是很优雅的。有了公式(BP1)和(BP2)，就可以计算出所有层对加权输入的误差了。 ${\delta }^L\to {\delta }^{L-1}\to \cdots \to {\delta }^l\to \cdots \to {\delta }^2\to {\delta }^1$ 。这也就是称之为梯度反向传播算法的原因。

证明：

从定义出发，应用多元微积分的链式法则：有 ${\delta }_j^l=\frac{\partial C}{\partial z_j^l}$ ，以及 ${\delta }_j^{l+1}=\frac{\partial C}{\partial z_j^{l+1}}$ 。 $\begin{array}{rl}{\delta }_j^l& =\frac{\partial C}{\partial z_j^l}\\ & =\sum _k\frac{\partial C}{\partial z_k^{l+1}}\frac{\partial z_k^{l+1}}{\partial z_j^l}\\ & =\sum _k\frac{\partial z_k^{l+1}}{\partial z_j^l}{\delta }_j^{l+1}\end{array}$

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

$z_k^{l+1}=\sum _j{w}_{kj}^{l+1}{a}_j^l+b_k^{l+1}=\sum _j{w}_{kj}^{l+1}\sigma (z_j^l)+b_k^{l+1}$

取微分，得到：

$\frac{\partial z_k^{l+1}}{\partial z_j^l}=w_{kj}^{l+1}{\sigma }^{{}^{\prime }}(z_j^l)$

带入上上式中，得到： ${\delta }_j^l=\sum _k{w}_{kj}^{l+1}{\delta }_k^{l+1}{\sigma }^{{}^{\prime }}(z_j^l)$

重写为矩阵形式，即：

${\delta }^l=((w^{l+1}{)}^T{\delta }^{l+1})\odot {\sigma }^{{}^{\prime }}(z^l)$

2.2.3 偏置的梯度\frac{\partial C}{\partial b^l_j}=\delta^{l}_{j}\qquad\qquad\qquad \cdots\cdots(BP3) \

第$l$层偏置的梯度就等于第$l$层加权输入的误差 ${\delta }^l$ 。改写为矩阵形式：

$\frac{\partial C}{\partial b^l}={\delta }^l$

证明：

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

$\frac{\partial C}{\partial b_j^l}=\frac{\partial C}{\partial z_j^l}\frac{\partial z_j^l}{\partial b_j^l}={\delta }_j^l\frac{\partial z_j^l}{\partial b_j^l}$

注意到： $z_j^l=\sum _k{w}_{jk}^l{a}_k^l+b_j^l$ ，故 $\frac{\partial z_j^l}{\partial b_j^l}=1$ 。带入上式中，得到：

\frac{\partial C}{\partial b^l_j}=\delta^{l}_{j}\2.2.4 权重的梯度

$\frac{\partial C}{\partial w_{jk}^l}=a_k^{l-1}{\delta }_j^l\cdots \cdots (BP4)$ 第$l$层权重的梯度就等于第$l$层加权输入的误差 ${\delta }_j^l$ 与 上一层神经元输出激活 $a_k^{l-1}$ 的乘积。将上式改写为：

$\frac{\partial C}{\partial w}=a_{in}{\delta }_{out}$

假设权重 $w$ 连接了两个神经元， $a_{in}$ 是神经元输入给权重的激活， ${\delta }_{out}$ 是权重输出给神经元的误差。

证明：

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

$\frac{\partial C}{\partial w_{jk}^l}=\frac{\partial C}{\partial z_j^l}\frac{\partial z_j^l}{\partial w_{jk}^l}={\delta }_j^l\frac{\partial z_j^l}{\partial w_{jk}^l}$ 注意到： $z_j^l=\sum _k{w}_{jk}^l{a}_k^l+b_j^l$ ，故 $\frac{\partial z_j^l}{\partial w_{jk}^l}=a_k^{l-1}$ 。带入上式中，得到：

$\frac{\partial C}{\partial w_{jk}^l}=a_k^{l-1}{\delta }_j^l$

2.2.5 四个基本等式

综上，梯度反向传播的四个基本等式是：

$\begin{array}{rl}{\delta }^L& ={\mathrm{\nabla }}_{a}C\odot {\sigma }^{{}^{\prime }}(z^L)\\ {\delta }^l& =(w^{l+1}{)}^T{\delta }^{l+1}\odot {\sigma }^{{}^{\prime }}(z^l)\\ \frac{\partial C}{\partial b_j^l}& ={\delta }_j^l\\ \frac{\partial C}{\partial w_{jk}^l}& =a_k^{l-1}{\delta }_j^l\end{array}$ 有了这四个基本等式，我们就基本理解了梯度的计算过程。

3. 后向传递与前向传递的FLOPs比率

3.1 定义

FLOPS：全大写，floating point operations per second，每秒钟浮点数计算次数，理解为运算速度，是衡量硬件性能的指标。例如：A100的float32运算速度是9.7TFLOPS，V100的float32运算速度是7.8TFLOPS。

FLOPs：s为小写，floating point operations，表示浮点数运算次数，理解为计算量，衡量模型或算法的复杂度。

如何计算矩阵乘法的FLOPs呢？ 对于 $A\in R^{1\times n},B\in R^{n\times 1}$ ，计算 $AB$ 需要进行 $n$ 次乘法运算和 $n$ 次加法运算，共计 $2n$ 次浮点数运算，需要 $2n$ FLOPs。对于 $A\in R^{m\times n},B\in R^{n\times p}$ ，浮点数运算次数为 $2mnp$ 。

backward-forward FLOP ratio：后向传递与前向传递的FLOPs比率，表示一次后向传递的浮点数计算次数与一次前向传递的浮点数计算次数的比率。衡量了一次后向传递要比一次前向传递多进行的浮点数运算。

3.2 前向传递与后向传递的浮点数操作次数的理论分析

为了理解后向传递的浮点数运算为什么要比前向传递多？我们就要从反向传播算法的四个基本等式入手。

我们从一个有2层隐藏层的神经网络入手：

我们假设线性层采用ReLU激活函数，采用随机梯度下降优化器。下表中h1和h2分别表示第1层和第2层隐藏层的神经元个数。有2层隐藏层的神经网络一次训练迭代的计算过程如下表所示：

从上表中，我们可以观察到，对于多层前馈神经网络模型，有以下结论：

1. 相比于线性层，激活函数ReLU和损失函数的浮点数运算量可以忽略。

2. 对于第一层，后向传递-前向传递的FLOPs比率是1:1。

3. 对于其他层，后向传递-前向传递的FLOPs比率是2:1。

4. 采用随机梯度下降作为优化器，权重更新的FLOPs是模型参数规模的2倍。

3.2.1 第一层与其他层的区别

对于多层前馈神经网络模型，采用随机梯度下降作为优化器。

第一层的后向传递-前向传递的FLOPs比率是1:1，其他层的后向传递-前向传递的FLOPs比率是2:1。模型参数更新的FLOPs是模型参数规模的2倍。

3.2.2 batch_size的影响

前向传递和后向传递的计算量FLOPs与batch_size成正比，即随着batch_size增大而线性增长。

优化器更新模型参数的FLOPs与batch_size无关，只与模型参数规模和优化器类型有关。

随着batch_size增大，前向传递和后向传递的FLOPs线性增长，而权重更新的FLOPs保持不变，这使得权重更新的FLOPs变得逐渐可以忽略不计了。

3.2.3 网络深度的影响

神经网络层数对后向传递-前向传递FLOPs比率有着间接的影响。由于第一层的后向-前向FLOPs比率是1:1，而其他层后向-前向FLOPs比率是2:1。层数的影响实际上是第一层与其他层的比较。

定义一个CNN神经网络，中间层的数量由0到100，可以看到后向-前向FLOPs比率由1.5提高到了接近2的水平。当层数逐渐变深的时候，第一层对模型整体的FLOPs影响就变小了，模型整体的后向-前向FLOPs比率就很接近2。

4. 总结

随着batch_size的增大，前向传递和后向传递的FLOPs线性增长，而权重更新的FLOPs保持不变，参数更新的FLOPs变得可以忽略不计了。这体现为：当batch_size足够大时，在训练神经网络的一次迭代中，前向传递和后向传递是主要的耗时环节，而参数更新的耗时变得几乎可以忽略不计。

由于第一层的后向-前向FLOPs比率是1:1，而其他层后向-前向FLOPs比率是2:1。随着网络层数的加深，第一层对整体FLOPs的影响变得可以忽略不计了。这体现为：当网络层数足够深时，后向传递的耗时几乎是前向传递耗时的2倍。

总的来说，对于用大batch_size的深层神经网络来说，后向传递-前向传递的FLOPs比率接近于2:1，换句话说，后向传递的计算量大约是前向传递的两倍。

5. 参考链接

1. Li S, Zhao Y, Varma R, et al. Pytorch distributed: Experiences on accelerating data parallel training[J]. arXiv preprint arXiv:2006.15704, 2020.

2. Narayanan D, Shoeybi M, Casper J, et al. Efficient large-scale language model training on gpu clusters using megatron-lm[C]//Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. 2021: 1-15.

3. http://neuralnetworksanddeeplearning.com/chap2.html (opens new window)

4. https://www.alignmentforum.org/posts/fnjKpBoWJXcSDwhZk/what-s-the-backward-forward-flop-ratio-for-neural-networks (opens new window)