深度学习中的矩阵求导基础#

Denominator Layout and Numerator Layout#

在深度学习中，矩阵求导是一个重要的概念。我们通常使用两种布局来表示矩阵求导：分母布局（Denominator Layout）和分子布局（Numerator Layout）。

分母布局（Denominator Layout） 在分母布局中，我们将矩阵的每一行视为一个整体，求导时将整个行作为一个单位来处理。这种布局在处理一些特定类型的矩阵时非常有用，例如在计算雅可比矩阵（矩阵的一阶导，Jacobian Matrix）时。

分子布局（Numerator Layout） 在分子布局中，我们将矩阵的每一列视为一个整体，求导时将整个列作为一个单位来处理。这种布局在处理一些其他类型的矩阵时更为方便，例如在计算海森矩阵（矩阵的二阶导，Hessian Matrix）时。

对于一个函数 $\mathbf{y} = \mathbf{W}\mathbf{x}$ ，其中 $\mathbf{y}\in\mathbb{R}^{D_1}, \mathbf{W}\in\mathbb{R}^{D_1 \times D_2}, \mathbf{x}\in\mathbb{R}^{D_2}$ ，其导数为:

$$\frac{\partial \mathbf{y}}{\partial \mathbf{x}} = \left[\begin{array}{cccc} \frac{\partial y_1}{\partial x_1} & \frac{\partial y_2}{\partial x_1} & \cdots & \frac{\partial y_{D_1}}{\partial x_1} \\ \frac{\partial y_1}{\partial x_2} & \frac{\partial y_2}{\partial x_2} & \cdots & \frac{\partial y_{D_1}}{\partial x_2} \\ \vdots & \vdots & \ddots & \vdots \\ \frac{\partial y_1}{\partial x_{D_2}} & \frac{\partial y_2}{\partial x_{D_2}} & \cdots & \frac{\partial y_{D_1}}{\partial x_{D_2}} \end{array}\right] = \mathbf{W}^\top \textrm{ (Denominator Layout)}$$ $$\frac{\partial \mathbf{y}}{\partial \mathbf{x}} = \left[\begin{array}{cccc} \frac{\partial y_1}{\partial x_1} & \frac{\partial y_1}{\partial x_2} & \cdots & \frac{\partial y_1}{\partial x_{D_2}} \\ \frac{\partial y_2}{\partial x_1} & \frac{\partial y_2}{\partial x_2} & \cdots & \frac{\partial y_2}{\partial x_{D_2}} \\ \vdots & \vdots & \ddots & \vdots \\ \frac{\partial y_{D_1}}{\partial x_1} & \frac{\partial y_{D_1}}{\partial x_2} & \cdots & \frac{\partial y_{D_1}}{\partial x_{D_2}} \end{array}\right] = \mathbf{W} \textrm{ (Numerator Layout)}$$

矩阵求导的链式法则#

对于一个函数 $y=f(g(h(x)))$ ，如果我们令 $q=h(x)$ ，$k=g(q)$ ，$y=f(k)$ ，那么根据链式法则，我们先以内层为单位排列（分母表达式，Denominator Layout）可以得到：

$$\frac{\partial y}{\partial x} = \frac{\partial q}{\partial x} \cdot \frac{\partial k}{\partial q} \cdot \frac{\partial y}{\partial k}$$

分子表达式（Numerator Layout）下：

$$\frac{\partial y}{\partial x} = \frac{\partial y}{\partial k} \cdot \frac{\partial k}{\partial q} \cdot \frac{\partial q}{\partial x}$$

一个例子#

对于函数 $y = (\mathbf{A}\mathbf{x}-\mathbf{b})^\top(\mathbf{A}\mathbf{x}-\mathbf{b})$ ，我们可以使用链式法则来计算其导数。

其中 $\mathbf{x}\in\mathbb{R}^{D}$，$\mathbf{A}\in\mathbb{R}^{D_1 \times D}$ ，$\mathbf{b}\in\mathbb{R}^{D_1}$ ，$y \in \mathbb{R}$ 。

令 $\mathbf{z} = \mathbf{A}\mathbf{x} - \mathbf{b}$ ，则有 $y = \mathbf{z}^\top \mathbf{z}$ 。

根据链式法则，我们可以得到：

$$\textrm{ (Denominator Layout)} \\ \begin{aligned} \frac{\partial y}{\partial \mathbf{x}} & =\frac{\partial \mathbf{z}}{\partial \mathbf{x}}\cdot\frac{\partial y}{\partial \mathbf{z}} \\ & =\mathbf{A}^\top\cdot(2\mathbf{z}) \\ & =2\mathbf{A}^\top(\mathbf{A}\mathbf{x}-\mathbf{b}) \end{aligned}$$ $$\textrm{ (Numerator Layout)} \\ \begin{aligned} \frac{\partial y}{\partial \mathbf{x}} & =\frac{\partial y}{\partial \mathbf{z}}\cdot\frac{\partial \mathbf{z}}{\partial \mathbf{x}} \\ & =2\mathbf{z}^\top\cdot\mathbf{A} \\ & =2(\mathbf{A}\mathbf{x}-\mathbf{b})^\top\cdot\mathbf{A} \end{aligned}$$

显然 $(2(\mathbf{A}\mathbf{x}-\mathbf{b})^\top\cdot\mathbf{A})\top = 2\mathbf{A}^\top(\mathbf{A}\mathbf{x}-\mathbf{b})$ 。

在分母表达式下 $2\mathbf{A}^\top(\mathbf{A}\mathbf{x}-\mathbf{b})$ 的形状是 $2\times \mathbb{R}^{D \times D_1} \times (\mathbb{R}^{D_1 \times D} \times \mathbb{R}^{D \times 1} - \mathbb{R}^{D_1 \times 1}) = \mathbb{R}^{D_1 \times 1}$ ，即为长度为 $D_1$ 的列向量。

含有Layer Norm的矩阵求导#

在深度学习中，Layer Norm（层归一化）是一个常用的技术，它在特征维度上进行归一化处理。

先假设输入是一个矩阵 $\mathbf{X} \in \mathbb{R}^{m \times n}$ ，其中 $m$ 是样本数量，$n$ 是特征数量。

1. 对于每个样本，计算每个样本的均值和标准差：
   • 均值：$\mu_i = \frac{1}{n} \sum_{j=1}^{n} x_{ij}$
   • 标准差：$\sigma_i^2 = \frac{1}{n} \sum_{j=1}^{n} (x_{ij} - \mu_i)^2$
2. 对每个样本进行归一化处理：
   • 归一化：$\hat{x}_{ij} = \frac{x_{ij} - \mu_i}{\sqrt{\sigma_i^2 + \epsilon}}$
   • 其中$\epsilon$是一个小常数，用于防止除以零。
3. 对归一化后的结果进行缩放和平移：
   • 平移：$y_{ij} = \gamma \hat{x}_{ij} + \beta$
   • 其中 $\gamma$ 和 $\beta$ 是可学习的参数。

例子：#

输入 $x$ 是一维列向量，长度为 $D$

$$\begin{aligned} & p=Wx \\ & k=\max(p,0) \\ & y=LN(k)=\gamma\odot\frac{k-\mu(k)}{\sqrt{\sigma^2(k)+\xi}}+\beta \end{aligned}$$

我们可以知道对于第一个式子 $\frac{\partial p}{\partial x} = W^\top$ 。

对于第二个式子p，它是一个列向量，它的值为 $\max(p,0)$ ，因此其导数为：

$$\frac{\partial k}{\partial p} = \begin{cases} 1 & \text{if } p > 0 \\ 0 & \text{if } p \leq 0 \end{cases}$$

一般我们可以将其表示为一个对角矩阵 $diag(\mathbf{1}_{p>0})$，其中对角线上的元素为 $1$ 或 $0$，具体取决于 $p$ 的值。

第三个式子就是重头戏了，Layer Norm的导数计算相对复杂，我们假设输入是一个列向量 $k$ ，其均值为 $\mu(k)$ ，标准差为 $\sigma^2(k)$ ，则Layer Norm的导数可以表示为：

$$k - \mu(k) = (\mathbf{I} - \frac{1}{D}\mathbf{1}\mathbf{1}^\top)k$$

其中 $\mathbf{1}\mathbf{1}^\top$ 是一个全1的矩阵。$D$的大小是$k$的长度。

我们令 $k - \mu(k) = \nu$ ，则有：

$$\begin{aligned} q = \frac{k-\mu(k)}{\sqrt{\sigma^2(k)+\epsilon}} & =\frac{k-\mu(k)}{\sqrt{\frac{1}{D}\sum_{i=1}^D(k_i-\mu(k))(k_i-\mu(k))+\xi}} \\ & =\frac{\nu}{\sqrt{\frac{1}{D}\sum_{i=1}^{D}\nu_i\nu_i+\xi}} \\ & =\frac{\sqrt{D}\nu}{\sqrt{||\nu||_2^2+D\xi}} \end{aligned}$$

那么 $\frac{\partial q}{\partial \nu}$ 的雅可比矩阵为：

$$\frac{\partial q}{\partial \nu} = \frac{\sqrt{D}}{\sqrt{||\nu||_2^2+D\xi}} \cdot (\mathbf{I} - \frac{\nu\nu^\top}{||\nu||_2^2+D\xi})$$

对于 $\mathbf{I} - \frac{\nu\nu^\top}{||\nu||_2^2+D\xi}$我们可以知道如果输入是 $\mathbb{R}^D$ ，那么输出就是 $\mathbb{R}^{D \times D}$ 。而它前面的部分就是一个标量了。

如果输入 $k$ 的每一维都相等，则$\nu$的所有分量为零，此时 $||\nu||_2^2$ 为零。为了避免这种情况导致分母为零，我们引入了 $\xi$ 作为正则化项，确保计算的稳定性。不加 $\xi$ 的话它就不是Lipschitz连续函数。关于Lipschitz连续后面会单独写一篇文章介绍。

回到 $y=LN(k)=\gamma\odot\frac{k-\mu(k)}{\sqrt{\sigma^2(k)+\xi}}+\beta$ ，可以重写为：

$$y=diag(\gamma)\cdot q+\beta$$

根据链式法则，我们可以得到：

$$\begin{aligned} \frac{\partial y}{\partial k} & = \frac{\partial \nu}{\partial k}\cdot\frac{\partial q}{\partial \nu}\cdot\frac{\partial y}{\partial q} \\ & = (\mathbf{I} - \frac{1}{D}\mathbf{1}\mathbf{1}^\top) \cdot \frac{\sqrt{D}}{\sqrt{||\nu||_2^2+D\xi}} \cdot (\mathbf{I} - \frac{\nu\nu^\top}{||\nu||_2^2+D\xi}) \cdot diag(\gamma) \end{aligned}$$

最终的$\frac{\partial y}{\partial x}$就是：

$$\frac{\partial y}{\partial x} = \frac{\partial p}{\partial x} \cdot \frac{\partial k}{\partial p} \cdot \frac{\partial y}{\partial k}$$

带入相应结果即可。

含有Loss#

我们知道神经网络需要计算损失函数（Loss Function）来进行训练。

对于下图所示的神经网络，我们可以使用链式法则来计算损失函数对输入的导数。 该神经网络有两个linear层，紧跟其后的是Softmax层和Cross Entropy Loss层。

在前向计算过程中我们要计算

$$\begin{aligned} & \mathbf{y} = W_1 \mathbf{x} \\ & \mathbf{h} = W_2 \mathbf{y} \\ & \mathbf{p} = \textrm{Softmax}(\mathbf{h}) \\ & \mathcal{L} = \textrm{CrossEntropy}(\mathbf{p}, \mathbf{t}) \end{aligned}$$

Softmax和CrossEntropy没有参数，实际上我们只需要计算两个线性层的导数。

注意这里的$\mathbf{t}$是一个one-hot编码的向量，表示真实标签。对于Softmax层和CrossEntropy层导数推导如下：

Softmax函数定义为：$\mathbf{p}_i = \frac{e^{\mathbf{h}_i}}{\sum_{j} e^{\mathbf{h}_j}}$ 导数为：$\frac{\partial \mathbf{p}}{\partial \mathbf{h}} = \textrm{diag}(\mathbf{p}) - \mathbf{p} \mathbf{p}^\top$ $\textrm{diag}(\mathbf{p})$ 是一个对角矩阵，主对角线是 $(\mathbf{p})$ 的元素。

CrossEntropy Loss定义为：$\mathcal{L} = -\sum_{i} \mathbf{t}_i \log(\mathbf{p}_i)$ 导数为：$\frac{\partial \mathcal{L}}{\partial \mathbf{p}} = -\frac{\mathbf{t}}{\mathbf{p}}$

根据链式法则，我们可以得到：$\frac{\partial \mathcal{L}}{\partial \mathbf{h}} = \mathbf{p} - \mathbf{t}$

$$\begin{aligned} & \frac{\partial \mathcal{L}}{\partial \mathbf{h}} = \frac{\partial \mathbf{p}}{\partial \mathbf{h}} \cdot \frac{\partial \mathcal{L}}{\partial \mathbf{p}} = \underbrace{\mathbf{p} - \mathbf{t}}_{\mathbb{R^C}} \\ & \frac{\partial \mathcal{L}}{\partial \mathbf{W_2}} = \underbrace{\frac{\partial \mathcal{L}}{\partial \mathbf{h}}}_{\mathbb{R^{C \times 1}}} \cdot \underbrace{\mathbf{y}^\top}_{\mathbb{R^{1 \times D_1}}} \\ & \frac{\partial \mathcal{L}}{\partial \mathbf{y}} = \underbrace{W_2^\top}_{\mathbb{R^{D_1 \times C}}} \cdot \underbrace{\frac{\partial \mathcal{L}}{\partial \mathbf{h}}}_{\mathbb{R^{C \times 1}}} \\ & \frac{\partial \mathcal{L}}{\partial \mathbf{W_1}} = \frac{\partial \mathcal{L}}{\partial \mathbf{y}} \cdot \mathbf{x}^\top \end{aligned}$$

参数更新#

TODO