1. Autograd (1)：PyTorch 自动一阶求导在标量、向量、矩阵、张量运算中的定义#
创建时间：2019-12-18
这一份文档中，我们将尝试理解 PyTorch 的自动求导的一些结论。
我们将不深究 PyTorch 的求导过程与程序问题，比如叶节点或中间矩阵导数等。我们只是探究不同的矩阵运算下自动求导的结论。
同时，我们也只讨论一阶导数的自动求导。一般的机器学习任务也只关心一阶导数。但在譬如分子力场等学习目标就包含了一阶导数的应用中，二阶导数可能是有意义的。二阶或高阶函数的求导在 PyTorch 中是存在的 (参考 Stackoverflow 问答 [1])，关于这部分探讨可能会放在以后的文档。
import torch
import numpy as np
torch.random.manual_seed(0)
torch.set_printoptions(precision=5, sci_mode=False, linewidth=120)
默认签名更改
这篇文档中，由于我们会经常地在求取自变量梯度后再次使用 backward 函数；为了避免程序的复杂性，我们会将 retain_graph 可选参数设为 True。
torch.Tensor.backward
<function torch.tensor.Tensor.backward(self, gradient=None, retain_graph=None, create_graph=False)>
torch.Tensor.backward.__defaults__ = (None, True, False)
1.1. 一元函数的自动求导#
谈及导数，最容易想到的是一元函数的导数。举例来说现在我们定义函数
当 \(b = 3\) 时，\(y \simeq 31.07\)。
b = torch.tensor(3., requires_grad=True)
y = b**3 + 10 * torch.exp(-b**2 / 10)
y
tensor(31.06570, grad_fn=<AddBackward0>)
利用求导法则，我们会很容易地知道，
当 \(b = 3\) 时，\(y^b \simeq 24.56\)。
符号定义
这篇文档中，定义 \(b, \boldsymbol{b}, \textbf{B}\) 分别为标量、向量、多维张量自变量；\(\boldsymbol{a}, \textbf{A}\) 分别为向量、矩阵常量。\(y, \boldsymbol{y}, \textbf{Y}\) 为表因变量 (函数)，\(z\) 为实际的标量因变量 (函数)。
作为张量分量的下角标统一采用 \(i, j, k, l, \cdots\)。
在这篇文档中，导数用类似于上述上标 \(y^b\) 的方式来简写定义。这样的定义方式可能是不常规或不合适的，但在后面定义矩阵导数、查看矩阵维度时多少会有方便之处。
float(3 * b**2 - 2 * b * torch.exp(-b**2 / 10))
24.56058120727539
写出上面一行代码显然对我们不仅仅有着脚本技工的要求，还有着初级的微积分的能力；除此之外，这种代码的可移植性差，若是换一个函数就需要重写一行导数代码。
当然，我们也可以用数值的方式求取导数。这样至少能保证代码的可移植性：
def num_deriv_3p_stencil(func, var, interval):
return (func(var + interval) - func(var - interval)) / (2 * interval)
y_func = lambda b: b**3 + 10 * np.exp(-b**2 / 10)
num_deriv_3p_stencil(y_func, 3, 1e-6)
24.560582046362356
但 PyTorch 作为包含了自动求导这个强大工具的程序库，我们不一定需要手动或数值地求取导数就可以给出解析的导数结果：
b.grad = None
y.backward()
b.grad
tensor(24.56058)
上面代码输出的精度可能不算太高。如果我们显示更多小数位数，我们会发现其精度比数值导数的精度高出很多：
float(b.grad)
24.56058120727539
关于数值导数的精度，读者可以尝试调整各种 interval 的数值；但不论多少 interval，误差总在 1e-7 或更高。
关于自动求导，需要补充的是，其中的 y.backward 允许引入一个与 y 相同维度的张量 (在这个例子中，y 是一个标量)。这个张量在函数的签名中是 gradient，其意义相当于是规定作为参数 \(b\) 的导数方向性。
拿现在具体的例子而言，我们认为 \(y\) 变量也是具有导数的，其导数定义为
其中，\(z\) 是一个标量，它代表真正用于计算的损失函数。尽管我们也称 \(y\) 是损失函数；若梯度 \(y^b = 0\) 时损失函数也确实降到了最低值 (若损失函数是凸的)；但实际程序中并不是拿 \(y^b\) 计算，而是使用
来作为真正的导数 b.grad。之所以上面我们安全地说 \(y_b \simeq 24.56\) 是因为在 PyTorch 中，\(y\) 维度为 1 时的 \(z^y\) 默认值就是 1。但若我们打算自己定义 \(z^y = 10\)，那么可以预期给出的 b.grad 的值为
b.grad = None
gy = torch.tensor(10.)
y.backward(gradient=gy)
b.grad
tensor(245.60582)
定义 gy \(z^y\) 看起来似乎不是很重要，在实际机器学习的训练中这与定义学习率的效果应当会是等价的。但在下述两个问题中，定义 \(z^y\) 可能有着它的意义：
\(z^y\) 若恰好是 \(y^{bb} = \frac{\partial^2 y}{\partial b^2}\)，那么 \(z^b\) 可以看作是凸优化区域中牛顿法数值迭代的 \(b\) 的方向。这应当比单纯的梯度下降的效率高不少。关于这一点的说明与应用可能会在以后的二阶自动求导的文档中给出。
若 \(y, b\) 不再是一元的标量，那么从程序的角度出发必须要求 \(z^y\) 的定义。关于这一点后文马上就会说明。
1.2. 向量点积的自动求导#
现在我们将 \(\boldsymbol{a}, \boldsymbol{b}\) 当作向量。
b = torch.randn(3, requires_grad=True)
b
tensor([ 1.54100, -0.29343, -2.17879], requires_grad=True)
a = torch.randn(3)
a
tensor([ 0.56843, -1.08452, -1.39860])
我们定义向量点积 \(y = \boldsymbol{a} \boldsymbol{b}\) (由于这里不涉及矩阵，因此就没有明确写出转置)，或更清晰地，
其中，\(i\) 取值范围在 0-2 之间。
y = a @ b
y
tensor(4.24143, grad_fn=<DotBackward>)
向量点积的导数是非常容易给出的，这可以表现为向量表示
也可以表现为标量表示
y.backward()
b.grad
tensor([ 0.56843, -1.08452, -1.39860])
关于向量或矩阵的微分问题，可以参考 ^Matrix Calculus 网站的结果。譬如对于当前的 \(y = \boldsymbol{a} \boldsymbol{b}\) 关于向量 \(\boldsymbol{b}\) 的导数，可以使用下述的表达式：
derivative of
a' * bw.r.t.b
应当能很容易地发现其导数就是 \(\boldsymbol{a}\)。
上述的讨论中，我们始终假定了 \(z^y = 1\)。
1.3. 矩阵-向量乘积的自动求导#
下面我们讨论更为广泛的问题。对于矩阵 \(\mathbf{A}\) 与向量 \(\boldsymbol{b}\) 的乘积
或者更详细地，
A = torch.randn(4, 3)
A
tensor([[ 0.40335, 0.83803, -0.71926],
[-0.40334, -0.59664, 0.18204],
[-0.85667, 1.10060, -1.07119],
[ 0.12270, -0.56632, 0.37311]])
b = torch.randn(3, 1, requires_grad=True)
b
tensor([[-0.89200],
[-1.50911],
[ 0.37039]], requires_grad=True)
y = A @ b
y.retain_grad()
y
tensor([[-1.89086],
[ 1.32759],
[-1.29354],
[ 0.88338]], grad_fn=<MmBackward>)
这个时候我们可能会遇到两个理解上的困境：到底导数要如何计算？程序要如何实现？
关于导数要如何计算，根据 Matrix Calculus 网站的结果，
或者更详细地，
这样的结论是比较容易理解的。
但从程序的角度上，既然求自动导数的目的是给 \(\boldsymbol{b}\) 向量 (作为自变量) 以一个下降的方向，让 \(\boldsymbol{y}\) 向量 (作为因变量) 的值变小，那么这个梯度量 b.grad 应当与 \(\boldsymbol{b}\) b 的维度相等才是。显然，\(\mathbf{A}\) 从维度上就已经不可能是一个合理的 b.grad。
不仅如此，PyTorch 不允许不加假定地使用 y.backward：
try:
y.backward()
except RuntimeError as e:
print("\033[31mRuntimeError: \033[0m" + "".join(e.args))
RuntimeError: grad can be implicitly created only for scalar outputs
事实上，解决方案是需要引入关于 \(\boldsymbol{y}\) 的偏导数 \(z^{y_i} = \frac{\partial z}{\partial y_i}\)；需要注意 \(z\) 始终是一个标量。因此，真正作为 b.grad 的量并非是 \(y_i^{b_j}\)，而是 \(z^{b_j}\)：
但这就给我们了自由发挥的空间了。如果我们对 \(z^{y_i}\) 给出一个随机数组 gy：
gy = torch.randn_like(y)
我们应当可以验证，b.grad 作为 \(z^{b_j}\) 确实满足上述的表达式：
b.grad = None
y.backward(gy)
torch.allclose(b.grad, A.T @ gy)
True
但下一个问题会是，在通常的机器学习的任务中，gy \(z^{y_i}\) 一般应当要如何选取？
这里举的一个例子是 L1 范数误差。如果我们假设向量 \(\boldsymbol{y}\) 的目标 (target) 是零值，那么作为标量的误差函数 z \(z\) 可以写作
z = torch.nn.L1Loss()(A @ b, torch.zeros(A.shape[0], 1))
z
tensor(1.34885, grad_fn=<L1LossBackward>)
我们可以很容易地用程序对其求导：
b.grad = None
z.backward()
b.grad
tensor([[ 0.04317],
[-0.77540],
[ 0.58640]])
对于上式，可以知道此时 gy
gy = y.sign() / y.shape[0]
gy
tensor([[-0.25000],
[ 0.25000],
[-0.25000],
[ 0.25000]], grad_fn=<DivBackward0>)
我们可以用上述的 gy 给出与 torch.nn.L1Loss 一样的反向传播的效果：
b.grad = None
y.backward(gy)
b.grad
tensor([[ 0.04317],
[-0.77540],
[ 0.58640]])
关于 gy \(z^{y_i}\) 在程序实现上的意义就在这里表述结束了。
总结来说，狭义上的 \(y_i\) 对 \(b_j\) 的导数 \(y_i^{b_j}\) 应当是一个二维矩阵 (正如符号 \(y_i^{b_j}\) 具有双下标 \(i, j\) 这个事实)，且该矩阵为 \(A_{ij}\)。对于 PyTorch 的 Autograd 而言，一般来说，表达式 \(y_i^{b_j}\) 通常是不会求取的；而真正被求取的量是在给定 \(z\) 对 \(y_i\) 的导数 \(z^{y_i}\) 下，\(z\) 对 \(b_j\) 的导数 \(z^{b_j}\)。\(z^{b_j}\) 的导数与 \(b_j\) 作为向量的维度相同且向量长度为 \(\dim(j)\)。
尽管 \(y_i^{b_j} = A_{ij}\) 通常情况下是不会被直接求取的，但是否真的意味着无法被求取？答案是否定的。对于给定的整数 \(i_0\) 且满足 \(0 \leqslant i_0 < \dim(i)\)，我们若令 \(z^{y_i} = \delta_{ii_0}\)，那么
这意味着当 \(z^{y_i} = \delta_{ii_0}\) 时，我们计算得到的 b.grad \(z^{b_j}\) 会返回矩阵 \(\mathbf{A}\) 的第 \(i_0\) 行。如此往复，就可以通过 b.grad 的堆叠矩阵 gbs 反推出整个 \(\mathbf{A}\) 矩阵 A。
def gy_to_bgrad(gy):
b.grad = None
y.backward(gy)
return b.grad
gy_list = torch.eye(y.shape[0], dtype=torch.float)[:, :, None]
gbs = torch.stack([gy_to_bgrad(gy).flatten() for gy in gy_list])
gbs
tensor([[ 0.40335, 0.83803, -0.71926],
[-0.40334, -0.59664, 0.18204],
[-0.85667, 1.10060, -1.07119],
[ 0.12270, -0.56632, 0.37311]])
torch.allclose(gbs, A)
True
1.4. 任意张量的自动求导#
有了上面的讨论，我们可以尝试对任意张量的缩并过程进行导数求取，并与 PyTorch 的 Autograd 进行对比。
下面的例子是
其中，\(i, j, k, l, m\) 指代的维度分别为 \(3, 4, 5, 6, 7\)。
A = torch.randn(3, 4, 5, 6)
B = torch.randn(4, 5, 6, 7, requires_grad=True)
Y = torch.einsum("ijkl, jklm -> ilm", A, B)
Y.shape
torch.Size([3, 6, 7])
如果现在我们定义 \(z^{Y_{ilm}} = \frac{\partial z}{\partial Y_{ilm}}\) gY 为任意的与 \(Y_{ilm}\) 相同维度的张量；以及通过其进行 Autograd 导出的 \(z^{B_{jklm}} = \frac{\partial z}{\partial B_{jklm}}\) 为 gB：
gY = torch.randn_like(Y)
Y.backward(gY)
gB = B.grad
gB.shape
torch.Size([4, 5, 6, 7])
那么 \(z^{B_{jklm}}\) 可以展示为下述表达式：
torch.allclose(
torch.einsum("ijkl, ilm -> jklm", A, gY),
gB
)
True
尽管我们最后并没有拓展到任意的张量的自动求导过程；但上面的这个例子相信已经具有足够的代表性了。依据上述的例子，也应当可以很容易地退化到普通矩阵、向量乃至标量的自动求导问题。