⚠️ 注意

笔者假定读者已经具备一定的线性代数和微积分基础，熟悉矩阵运算和基本的求导规则。

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

矩阵求导的布局

在深度学习中，矩阵求导是一个重要的概念。我们通常使用两种布局来表示矩阵求导：分母布局（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} $$

从泰勒展开到梯度下降法和牛顿法

在深度学习中，参数更新是通过梯度下降（Gradient Descent）来实现的。我们使用计算得到的梯度来调整模型的参数，以最小化损失函数。

而在这之前，必须先来讲讲泰勒展开（Taylor Expansion）。

对于多元函数 $f(\mathbf{x}): \mathbb{R}^n \to \mathbb{R}$（例如损失函数），在点 $\mathbf{a}$ 处的泰勒展开为：

$$ f(\mathbf{x}) \approx f(\mathbf{x_0}) + \nabla f(\mathbf{x_0})^T(\mathbf{x}-\mathbf{x_0}) + \frac{1}{2}(\mathbf{x}-\mathbf{x_0})^T \mathbf{H}(\mathbf{x_0}) (\mathbf{x}-\mathbf{x_0}) + \cdots $$

其中，$\mathbf{H}$ 是海森矩阵（Hessian Matrix），表示二阶导数信息。

在深度学习中，我们通常只使用到一阶泰勒展开，即：

$$ f(\mathbf{x}) \approx f(\mathbf{x_0}) + \nabla f(\mathbf{x_0})^T(\mathbf{x}-\mathbf{x_0}) $$

通常优化条件下，我们希望找到一个 $\mathbf{x}$ 使得 $f(\mathbf{x})$ 最小化。根据一阶泰勒展开，此时 $f(\mathbf{x}) \leq f(\mathbf{x_0})$ 。那么我们只需要确保迭代之前的 $\nabla f(\mathbf{x_0})^T(\mathbf{x}-\mathbf{x_0})$ 项小于零即可。这又可以导出 $\nabla f(\mathbf{x_0})^T(\mathbf{x}-\mathbf{x_0}) < 0$ 。

那么就有 $\mathbf{x} - \mathbf{x_0} = -\alpha \nabla f(\mathbf{x_0})$ ，即

$$ \mathbf{x} = \mathbf{x_0} - \alpha \nabla f(\mathbf{x_0}) \quad \textrm{(Gradient Descent)} $$

最后会下降到梯度为零的点。那么 $\nabla f(\mathbf{x}) = 0$ 。

自然我们对上面的泰勒展开进行一次求导，得到：

$$ \nabla f(\mathbf{x}) \approx 0 + \nabla f(\mathbf{x_0}) + \mathbf{H}(\mathbf{x_0})(\mathbf{x}-\mathbf{x_0}) = 0 \\ \mathbf{H}(\mathbf{x_0})(\mathbf{x}-\mathbf{x_0}) = -\nabla f(\mathbf{x_0}) \\ \mathbf{x} = \mathbf{x_0} - \mathbf{H}(\mathbf{x_0})^{-1} \nabla f(\mathbf{x_0}) \quad \textrm{(Newton's Method)} $$

拓展部分

我们讲的内容都是为了接下来的自动微分服务。

现代的深度学习框架（如TensorFlow和PyTorch）已经实现了自动微分（Automatic Differentiation），使得我们可以轻松地计算复杂模型的梯度，而无需手动推导每个步骤。 自动微分通过构建计算图（Computation Graph）来跟踪所有的操作，并在反向传播（Backpropagation）过程中自动计算梯度。这大大简化了深度学习模型的训练过程，使得研究人员和工程师能够专注于模型设计和优化，而不必担心复杂的数学推导。

自动微分有前向模式（Forward Mode）和反向模式（Reverse Mode）两种主要实现方式。前向模式适用于输入维度较小的情况，而反向模式（常说的反向传播）则更适合深度学习中的大规模模型，因为它能够高效地计算标量输出相对于大量参数的梯度。