张量链式法则（上篇）：任意维度反向传播公式推导与常见算子解析

本文首发于本人的微信公众号，原文链接：https://mp.weixin.qq.com/s/k4hMPT4UrHoPcV3lcdTaiQ

摘要

本文为常用算子反向传播公式的上篇，介绍了适用于任意张量函数的链式法则公式，使用该公式可以求出诸如 reshape，broadcast_to 这类会改变张量维度数量的算子的反向传播公式。

本文同时给出了求常见算子反向传播公式的通用方法，并以几个简单的算子为例进行了演示。

本系列文章的下篇将用本文提到的公式求解 reshape，transpose 等更复杂的算子的反向传播公式。

写在前面

昨天写了一篇关于矩阵函数链式法则的文章，那篇文章里的方法只适用于二维矩阵函数的情形，且没有对公式的正确性做解释。

在经过了一天的思考和沉淀后，我发现了更加通用的链式法则公式，该公式适用于任意维度数量的张量函数（不要求输入输出的维度数量相同），并且还琢磨出了公式正确性的证明。

本文将会介绍这一公式，并使用这一公式来对12个常见的深度学习算子进行反向传播公式推导。

本文涉及到的算子全部来自CMU 10-714的Homework 1。

由于公式推导部分内容较多，所以这篇文章会分为上下两篇。

本篇主要介绍基本方法，以及一些简单算子的公式推导，涉及到维度数量变化的复杂算子公式推导会放到下篇单独展开。

张量的链式法则公式

上一篇文章里介绍的公式只适用于二维矩阵映射到二维矩阵的函数，但是在深度学习中通常还会遇到高维张量，甚至于遇到会改变张量维度数量的情况，例如 reshape，broadcast_to，summation。

本节介绍的张量链式法则公式可以适用于任意维度数量的张量函数，真正实现梯度自由。

问题描述

• \(X \in \mathbb{R}^{d_1 \times d_2 \times \cdots \times d_n}\) 是一个 \(n\) 维张量

• 函数 \(F: \mathbb{R}^{d_1 \times d_2 \times \cdots \times d_n} \to \mathbb{R}^{e_1 \times e_2 \times \cdots \times e_m}\) 将 \(n\) 维张量 \(X\) 映射到 \(m\) 维张量

• 损失函数 \(L: \mathbb{R}^{e_1 \times e_2 \times \cdots \times e_m} \to \mathbb{R}\) 将 \(m\) 维张量 \(F(X)\) 映射到一个实数，即 \(Y = L(F(X))\)

• 现在已知 \(G = \frac{\partial Y}{\partial F}\)

• 求 \(\nabla = \frac{\partial Y}{\partial X}\)

解决方法

首先先确定这里涉及到的张量的形状，可以断言 \(G \in \mathbb{R}^{e_1 \times e_2 \times \cdots \times e_m}\)，因为 \(L\) 是一个将 \(m\) 维张量映射到实数的函数，所以其导数也一定是和 \(Y\) 同形状的。

同理也可以得出 \(\nabla \in \mathbb{R}^{d_1 \times d_2 \times \cdots \times d_n}\)

记张量 \(G\) 的第 \(\mu_1, \mu_2, \ldots, \mu_m\) 个元素为 \(g_{\mu_1 \mu_2 \mu_3 \cdots \mu_m}\)

张量 \(F(X)\) 的第 \(\mu_1, \mu_2, \ldots, \mu_m\) 个元素为 \(f_{\mu_1 \mu_2 \mu_3 \cdots \mu_m}\)

最终梯度 \(\nabla\) 的第 \(\lambda_1, \lambda_2, \ldots, \lambda_n\) 个元素为 \(\nabla_{\lambda_1 \lambda_2 \lambda_3 \cdots \lambda_n}\)

张量 \(X\) 的第 \(\lambda_1, \lambda_2, \ldots, \lambda_n\) 个元素为 \(x_{\lambda_1 \lambda_2 \lambda_3 \cdots \lambda_n}\)

则有如下公式：

\[\nabla_{\lambda_1 \lambda_2 \cdots \lambda_n}=\sum_{\substack{\mu_1 \in [1, e_1] \\ \mu_2 \in [1, e_2] \\ \vdots \\ \mu_m \in [1, e_m]}} g_{\mu_1 \mu_2 \cdots \mu_m} \frac{\partial}{\partial x_{\lambda_1 \lambda_2 \cdots \lambda_n}} f_{\mu_1 \mu_2 \cdots \mu_m} \]

这个公式可以这样理解：

1. 我们首先把目光聚焦到某一个自变量 \(x_{\lambda_1 \lambda_2 \lambda_3 \cdots \lambda_n}\) 上
2. 然后用 \(F(X)\) 这个张量里的每个元素对 \(x_{\lambda_1 \lambda_2 \lambda_3 \cdots \lambda_n}\) 求导，得到一个 \(m\) 维张量，记为 \(F'\)
3. 然后再把这个张量 \(F'\) 和张量 \(G\) 做哈达马积得到 \(F''\)（即对应的元素相乘，这两个张量的形状是相同的）
4. 最后把 \(F''\) 里面的所有元素相加，得到一个实数，就是自变量 \(x_{\lambda_1 \lambda_2 \lambda_3 \cdots \lambda_n}\) 的导数，即 \(\nabla_{\lambda_1 \lambda_2 \lambda_3 \cdots \lambda_n}\)
5. 对 \(X\) 里的所有元素做同样的操作即可得到最终的导数 \(\nabla\)

正确性证明

上述操作其实是在做Tensor Contraction，但是为什么这样做就能求出正确的导数呢？下面我们来证明这个方法的正确性。

这里为了让文章看起来更简洁，我们就不失一般性地假设 \(X\) 和 \(F(X)\) 都是3维张量，且 \(X \in \mathbb{R}^{d_1 \times d_2 \times d_3}\)，\(F(X) \in \mathbb{R}^{e_1 \times e_2 \times e_3}\)，下面的推导过程可以很轻易地推广到任意维度数量的情形。

这里我们同样先把视角放到具体的某一个自变量 \(x_{ijk}\) 上，要记住一点，\(Y = L(F(X))\) 最终会把这个自变量映射到一个实数上，也就是说Y本质上是一个关于 \(x_{ijk}\) 的表达式

比如可能本质上 \(Y = x_{123} \cdot x_{345} + \cdots\)，我们完全可以使用简单的一元函数求导公式求出其导数。

但是目前的情况是，\(Y\) 实际上是以 \(F\) 来表达的，例如 \(Y = f_{123} + f_{456}\)，并且我们已知 \(F\) 的具体表达式（即 \(F\) 里的任一元素 \(f_{lmn}\) 是如何通过 \(X\) 里的元素计算出来的）。

那么这个时候就可以使用多元函数的链式法则，把 \(Y\) 看作是由 \(f_{111}, f_{112}, \ldots, f_{e_1 e_2 e_3}\) 构成的复合函数，那么就可以得到

\[\frac{\partial Y}{\partial x_{ijk}} = \sum_{\substack{l \in [1, e_1] \\ m \in [1, e_2] \\ n \in [1, e_3]}} \frac{\partial Y}{\partial f_{lmn}} \frac{\partial f_{lmn}}{\partial x_{ijk}} \]

其中 \(\frac{\partial Y}{\partial f_{lmn}}\) 就是 \(g_{lmn}\)，即张量 \(G\) 对应位置的元素

\(\frac{\partial f_{lmn}}{\partial x_{ijk}}\) 则是 \(F'\) 张量中的元素

求和符号里的整体就是 \(F''\) 里的元素

将这些元素全部相加就能够得到 \(Y\) 关于 \(x_{ijk}\) 的导数。

所以这里能够使用Tensor Contraction进行链式法则计算的本质原因是因为相乘再求和这一操作与多元函数链式法则是相一致的。

对于有多层嵌套的张量函数，例如神经网络，多次应用张量链式法则即可求得导数，因此这里只讨论 \(Y = L(F(X))\) 这一典型场景。

上述过程可以推广到任意维度数量的张量函数，具体过程不展开赘述了，留作习题自证不难。

对于损失函数是多维张量的场景（目前来看比较少见），也可以使用类似的过程进行推导，只是最终的导数结果会是一个多维张量，这里同样不展开赘述了。

常见算子的反向传播推导（上半部分）

基本方法

这里的问题和上文提到的张量函数链式法则问题相同，\(G\) 表示反向传播中上一个节点传递过来的导数。

基本的求解方法如下：

1. 首先确定涉及到的张量形状，尤其是函数 \(F\) 的输入张量与输出张量的形状（具体形状与函数定义，函数传入的参数相关）
2. 写出 \(f_{ijk \cdots}\) 的表达式
3. 然后根据上文提到的公式，写出 \(\nabla_{ijk \cdots}\) 的表达式，最后尝试化简表达式。

这里分享一些小技巧：

1. 对于输入输出形状不定的函数，例如与输入参数有关的 broadcast_to，可以先以一维向量 \(X\) 得到二维矩阵 \(F(X)\) 为例子进行推导，然后再尝试扩展到任意维度。
2. 在扩展时，可以尝试使用猜测的方式，因为最终导数的形状是已知的，所以可以尝试使用各种方法来凑出最终的形状，猜出表达式，然后再对表达式进行证明验证。

PowerScalar

这一算子的作用是将一个张量的所有元素都进行一个指定幂次的 pow 操作，以二维矩阵为例，\(f_{ij} = x_{ij}^{\text{scalar}}\)，其中 \(\text{scalar}\) 为算子的参数。

（注：算子的输入分为自变量和参数，这里求导仅针对自变量求导，参数可看作常数）

这里函数的输入和输出形状相同，且和参数无关，所以就直接假设 \(X\) 和 \(F(X)\) 为二维矩阵进行推导了，不难验证，推导出的表达式同样适用于任意维度数量的情形。

首先确定，如果 \(X\) 为 \(m \times n\) 矩阵，那么 \(F(X)\) 也是 \(m \times n\) 矩阵，所以 \(\nabla\) 也是 \(m \times n\) 矩阵，其拥有2个维度，所以首先需要写出 \(\nabla_{ij}\) 的表达式。

由于 \(F(X)\) 和 \(G\) 都是2个维度，所以求和变量有2个，分别求和到 \(m\) 和 \(n\)，所以最终的表达式为：

\[\nabla_{ij} = \sum_{k=1}^{m} \sum_{l=1}^{n} g_{kl} \cdot \frac{\partial f_{kl}}{\partial x_{ij}} = \sum_{k=1}^{m} \sum_{l=1}^{n} g_{kl} \cdot \frac{\partial x_{kl}^{\text{scalar}}}{\partial x_{ij}} \]

注意到，只有当 \(k = i\) 且 \(l = j\) 时，求和项才不为0，所以求和式化简后为

\[\nabla_{ij} = \text{scalar} \cdot g_{ij} \cdot x_{ij}^{\text{scalar} - 1} \]

所以最终求得导数

\[\nabla = \text{scalar} \cdot G \odot X^{\text{scalar} - 1} \]

注：这里的 \(X^a\) 是指 PowerScalar(X, a)，而非 \(a\) 个 \(X\) 矩阵进行矩阵乘法，后文也将沿用这一记法。

EWiseDiv

这一算子有两个参数 \(F(X, Y) = X / Y\)，这里的除法表示逐元素相除，即 \(f_{ij} = \frac{x_{ij}}{y_{ij}}\)，后文也将沿用这一记法。

这里同样以二维矩阵为例推导，首先确定如果 \(X\) 为 \(m \times n\) 矩阵，那么 \(\nabla\) 也是 \(m \times n\) 矩阵。

同样的，\(G\) 和 \(F(X, Y)\) 为 \(m \times n\) 矩阵，所以可以仿照上面 PowerScalar 的例子写出：

\[\nabla_{ij} = \sum_{k=1}^{m} \sum_{l=1}^{n} g_{kl} \cdot \frac{\partial}{\partial x_{ij}} \frac{x_{kl}}{y_{kl}} \]

注意到，只有当 \(k = i\), \(l = j\) 时表达式不为0，所以如果以 \(X\) 为自变量，可以得到

\[\nabla_{ij}^{X} = \frac{g_{ij}}{y_{ij}} \]

故

\[\nabla^{X} = \frac{G}{Y} \]

同理，可以得到对 \(Y\) 的导数为

\[\nabla^{Y} = -\frac{G \odot X}{Y^2} \]

DivScalar

这一算子是将 \(X\) 的所有元素都除以一个常数 \(\text{scalar}\)，具体推导过程和 PowerScalar 高度重合，这里不再赘述，只给出最终结果，推导过程留作习题供读者练习。

\[\nabla = \frac{G}{\text{scalar}} \]

（未完待续）