深度 | BP表达式与硬件架构：相似性构建更高效的计算单元

选自Medium

作者：Yaroslav Bulatov

机器之心编译

参与：蒋思源

反向传播是当前深度学习主要使用的参数更新方法，因此深度学习的硬件设计也需要拟合这种反向传播的计算结构。本文从反向传播的抽象表达开始简要地分析了 BP 算法和脉动阵列架构（systolic array architecture）之间的相似性，从而表明了脉动阵列架构适合执行 BP 和进行模型训练。

在并行计算的体系架构中，脉动阵列（systolic array）是紧密耦合的数据处理单元（data processing unit/DPU）的一种同构网络。每一个结点或 DPU 独立地计算部分结果，并将该部分结果作为从上游单元接受数据的函数，在将结果储存在当前结点后会传递到下游单元。本文重点在于论述反向传播算法的抽象表达，并讨论表达式与这种脉动阵列架构之间的相关性。

假设我们有一个 n=3 的多层网络架构，且运算的对象为 m 维向量，因此预测或推断的过程可以表示为：

我们将预测函数 Y 放入到损失函数 l 中以进行最优化。为了表示这种结构，我们使用圆圈表示复合函数算子（Ring 算子），因此目标函数 L 可以写为：

根据链式法则，目标函数的导数可以根据矩阵乘法的形式写为：

其中，链式乘法中的每一项都是雅可比矩阵（Jacobian matrix）。为了更形象地说明这一过程，假设我们的损失函数 l 有以下形式：

层级函数 f 仅仅只是简单地求输入向量中每一个元素的平方：

它们的雅可比矩阵就可以写为以下形式：

为了计算目标函数的导数，我们需要乘以这些雅可比矩阵。因此这种链式矩阵乘法的维度就可以可视化为以下形式：

计算这种链式矩阵乘法首先需要注意的就是确定乘法的顺序，因为矩阵的乘法分为左乘和右乘，且不同的乘法顺序会影响计算的效率，所以我们需要确定矩阵乘法的计算顺序。这种寻找最优乘法序列的任务称为矩阵链式排序问题。在本案例中，因为向量左乘矩阵还是得到一个向量，所以我们只需要从左往右进行矩阵乘积就能进行高效的计算。

其次我们需要考虑如何具体地计算这些矩阵运算而不使用构建雅可比矩阵。这是非常重要的，因为模型的特征数量 m 可能是几万的数量级，这意味着雅可比矩阵可能有数十亿的元素。在本案例中，雅可比矩阵是一个对角矩阵，那么向量和雅可比矩阵的乘积就等价于向量对应元素间的乘积，因此我们就能避免构建一个 m-x-m 的雅可比矩阵。

神经网络中典型的层级函数也采用这种高效实现的运算。这种向量-雅可比乘积（vector-Jacobian product）运算是任何反向传播算法实现的关键，Theano 称其为「Lop」（左乘算符）、PyTorch 称之为「backward」方法、TensorFlow 称之为「grad」或「grad_func」。

为了简化表达，我们将计算生成的中间值（即激活值）记为 A：

通过上图，我们将目标函数的导数写为：

因为损失函数的雅可比矩阵只是简单地转置输入矩阵，因此我们可以写为：

如果我们定义第一个反向传播的值为 B_3 = A_3'，那么我们可以将计算写为一个序列

为了进一步简化，令 b 指代向量-雅可比乘积（即 backwards()、Left operator、grad_func），使用 Hadamard 乘积的符号表示元素对应乘积。我们就可以将向量-雅可比乘积写为：

我们最终可以将前向/反向传播的公式写为：

这一过程的计算图可以表示为（以下两个计算图是等价的）：

如果我们查看二维 systolic array 的架构，就会发现它们之间的结构是非常相似的，也就是说这种硬件架构能很好地拟合反向传播算法。

原文链接：https://medium.com/@yaroslavvb/backprop-and-systolic-arrays-24e925d2050

本文为机器之心编译，转载请联系本公众号获得授权。