一种偏主观的矩阵乘法梯度推导方法

很早在纸上推导过梯度的计算方法，但每次都忘记推导过程反复推导。于此想总结新的记忆方法。

梯度下降推导过程难以记忆来自于矩阵微积分，矩阵微积分中涉及标量、向量、矩阵之间两两求导操作，其定义如下图，√ 表示存在定义，x 表示不存在定义^[1]：

函数 \ 自变量	scalar	vector	matrix
scalar	√	√(Nabla)	√
vector	√	√(Jacobi)	×
matrix	√	×	×

• 矩阵和矩阵之间、向量和矩阵之间不存在导数
• 导数的维度是函数、自变量维度之和，反应定义导数时函数和自变量之间的维度是正交关系，导数中每个元素都是标量求导

矩阵之间、矩阵向量之间不存在导数的原因是我们将定义范围限制在最高二维矩阵中，所以维度不能超过2，此类导数定义不存在。对于矩阵乘法 \(\mathbf{Y}_{B , O} = \mathbf{X}_{B , I} \mathbf{W}_{O , I}^T\) ，我们当然希望能够以矩阵形式优雅地推导整体公式，但矩阵-矩阵之间并不存在导数的定义，需要频繁在标量-向量-矩阵之间切换，或者引入许多额外 Nabla 在矩阵作用的关系，导致推理过程别扭不美观。

因此引入张量，允许变量维度超过 2 维，便可定义任意两个张量之间的导数。注意知道 \(\mathbf{Y}\) 的维度 B 来自 \(\mathbf{X}\)，维度 O 来自 \(\mathbf{W}\)，但从函数角度来看，只是定义了一个 \(R^{B\times I} , R^{O \times I} \rightarrow R^{B \times O}\) 的函数，应将这些维度用不同的符号区分。

定义矩阵乘法和维度：

\[\mathbf{Y}_{B' , O'} = \mathbf{X}_{B , I'} \mathbf{W}_{O , I}^T \]

已知损失函数 F 相对输出张量梯度：

\[\frac{\partial F}{\partial \mathbf{Y}} \in R^{B'\times O'} \]

则有

\[\begin{align*} \frac{\partial F}{\partial \mathbf{W}}&=\frac{\partial F}{\partial \mathbf{Y}} \frac{\partial \mathbf{Y}}{\partial \mathbf{W}} \in R^{O\times I}\\ \frac{\partial \mathbf{Y}}{\partial \mathbf{W}} &\in R^{B'\times O' \times O \times I} \end{align*} \]

这个过程中假设张量也满足链式法则，且链式法则传递关系通过张量乘法也就是 Einsum 规约相同符号维度，是否严格成立需要补充证明。

接下来需要求解具体 \(\frac{\partial \mathbf{Y}}{\partial \mathbf{W}}\) ，不幸的是仍然需要拆分到标量求导。这种定义方法只包含两个层次，多维度的张量表示，以及具体计算的标量表示。

\[\begin{align*} y_{b',o'} &= \sum_{i \in I} w_{o,i} x_{b,i'}\\ s.t.\quad b' &= b\\ s.t. \quad o' &= o\\ s.t. \quad i' &= i \end{align*} \]

可得：

\[\begin{cases} \frac{\partial y_{b',o'}}{\partial w_{o,i}} = x_{b',i},& o'=o\\ \frac{\partial y_{b',o'}}{\partial w_{o,i}} = 0,& o' \neq o \end{cases} \]

换句话说，\(\frac{\partial \mathbf{Y}}{\partial \mathbf{W}}\) 选取任意坐标 \((B=b', I=i)\) 做切片，切出来的 \(R^{O'\times O}\) 的矩阵是单位阵的倍数。

附带验证程序：

import torch
import einops

B = 8
O = 32
I = 256

B_ = B
O_ = O
I_ = I

x = torch.randn(B, I_).to("cuda")
w = torch.randn(O, I).to("cuda").requires_grad_(True)  
y = torch.einsum('bi,oi->bo', x, w)

df_dy = torch.randn(B_, O_).to("cuda")

with torch.no_grad():
    dy_dw = torch.zeros(B_, O_, O, I).to("cuda")
    for b_ in range(B_):
        print(f'batch {b_}')
        for o_ in range(O_):
            for o in range(O):
                for i in range(I):
                    if o_ == o:
                        dy_dw[b_, o_, o, i] = x[b_, i]

    df_dw = torch.einsum('BO,BOoi->oi', df_dy, dy_dw)

if w.grad is not None:
    w.grad.zero_()

y.backward(df_dy)  

auto_dy_dw = w.grad

if auto_dy_dw is None:
    print("Error: w.grad is None. Gradient was not properly calculated.")
else:
    diff = torch.abs(auto_dy_dw - df_dw)
    print(f'Max diff: {diff.max()}')

这边建模好处是用统一张量求导运算替代混乱的各种 矩阵-向量-标量 求导公式，但该求导方法普适和矩阵乘法无关，需要额外理解和建模维度之间的约束关系。到底哪种方便见仁见智了。

或者可以利用这种推导理解为什么直接用维度凑 einsum 表达式就可以直接表示梯度。

import torch
import einops

B = 8
O = 32
I = 256

x = torch.randn(B, I_).to("cuda")
w = torch.randn(O, I).to("cuda")
y = torch.einsum('bi,oi->bo', x, w)

df_dy = torch.randn(B, O).to("cuda")
df_dw = torch.einsum('bo,bi->oi', df_dy, x)

---

1. https://en.wikipedia.org/wiki/Matrix_calculus ↩︎