MindSpore多元自动微分

技术背景

当前主流的深度学习框架，除了能够便捷高效的搭建机器学习的模型之外，其自动并行和自动微分等功能还为其他领域的科学计算带来了模式的变革。本文我们将探索如何用MindSpore去实现一个多维的自动微分，并且得到该多元函数的雅可比矩阵。

函数形式与雅可比矩阵形式

首先我们给定一个比较简单的z关于自变量x的函数形式（其中y和I是一些参数）：

\[z_{i,j}(x)=y_ix_j \]

比如我们考虑一个3*3的z，我们最终需要计算的是这样一个雅可比矩阵：

\[J_z(x)= \left[ \begin{array}{l} \frac{\partial z_0}{\partial x_0} & \frac{\partial z_0}{\partial x_1} & \frac{\partial z_0}{\partial x_2}\\ \frac{\partial z_1}{\partial x_0} & \frac{\partial z_1}{\partial x_1} & \frac{\partial z_1}{\partial x_2}\\ \frac{\partial z_2}{\partial x_0} & \frac{\partial z_2}{\partial x_1} & \frac{\partial z_2}{\partial x_2} \end{array} \right] \]

假如我们给定一些简单的初始值：

\[x=[1,2,3]\\ y=[1,3,2] \]

那么理论上我们应该得到的结果是：

\[J_z(x)=\left[ \begin{array}{l} 1 & 0 & 0\\ 0 & 0 & 3\\ 0 & 2 & 0\\ \end{array} \right] \]

接下来我们看看如何在MindSpore的自动微分框架下实现这一功能。

初步尝试Grad自动微分

我们先按照上一章节中的公式的基本内容，直接写一个Net函数用于表示z，然后再用grad函数对其进行微分，代码内容如下所示：

from mindspore import nn, Tensor, ops
from mindspore.ops.functional import grad
import numpy as np
from mindspore import numpy as msnp

class Net(nn.Cell):
    def __init__(self, y, index):
        super(Net, self).__init__()
        self.y = y
        self.index = index
        self.norm = nn.Norm(-1)

    def construct(self, x):
        return self.y[self.index]*x

x = Tensor(np.array([1,2,3]).astype(np.float32))
y = Tensor(np.array([[1],[2],[3]]).astype(np.float32))
index = Tensor(np.array([0,2,1]).astype(np.int32))
shape = (y.shape[0], x.shape[0])

output = grad(Net(y,index))(x)
print(output)
# [6. 6. 6.]

在这个案例中，我们得到的结果，首先维度就不对，我们理想中的雅可比矩阵应该是3*3大小的，可见MindSpore中自动微分的逻辑是把其中的一个维度进行了加和，类似于这样的形式：

\[\left[ \frac{\partial z_0}{\partial x_0}+\frac{\partial z_1}{\partial x_0}+\frac{\partial z_2}{\partial x_0}, \frac{\partial z_0}{\partial x_1}+\frac{\partial z_1}{\partial x_1}+\frac{\partial z_2}{\partial x_1}, \frac{\partial z_0}{\partial x_2}+\frac{\partial z_1}{\partial x_2}+\frac{\partial z_2}{\partial x_2} \right] \]

所以为了得到我们的结果，需要对输入的x进行扩维。

尝试扩维输入的自动微分

在MindSpore中提供了BroadcastTo这样的接口，可以自动的在扩展维度填充待扩展张量的元素，我们需要把x的最外层维度扩展到与参数y一致，在这个案例中就是3*3的维度，具体代码实现如下所示：

from mindspore import nn, Tensor, ops
from mindspore.ops.functional import grad
import numpy as np
from mindspore import numpy as msnp

class Net(nn.Cell):
    def __init__(self, y, index):
        super(Net, self).__init__()
        self.y = y
        self.index = index
        self.norm = nn.Norm(-1)

    def construct(self, x):
        return self.y[self.index]*x

x = Tensor(np.array([1,2,3]).astype(np.float32))
y = Tensor(np.array([[1],[2],[3]]).astype(np.float32))
index = Tensor(np.array([0,2,1]).astype(np.int32))
shape = (y.shape[0], x.shape[0])

output = grad(Net(y,index))(ops.BroadcastTo(shape)(x))
print(output)
'''
[[1. 1. 1.]
 [3. 3. 3.]
 [2. 2. 2.]]
'''

从这个输出结果中我们发现，虽然维度上是被扩展成功了，但是那些本该为0的位置却出现了非0元素，这说明在自动微分计算的过程中，我们输入的参数y也被自动的Broadcast了，而实际上正确的计算过程中是不能使用Broadcast的。

为参数添加Mask

上一个章节中说道，如果利用Tensor本身的自动Broadcast会导致输入参数被扩维，会得到一个错误的微分结果。因此这里我们手动对输入参数进行正确的扩维，这个过程是添加一个Mask矩阵，用于标记每一个参数所对应的位置。这里我们假设输入一个这样的Mask矩阵：

\[I=\left[ \begin{array}{l} 1 & 0 & 0\\ 0 & 0 & 1\\ 0 & 1 & 0 \end{array} \right] \]

这样理论上最终微分结果的非0元素应该跟这个矩阵是一致的，相关代码如下所示：

from mindspore import nn, Tensor, ops
from mindspore.ops.functional import grad
import numpy as np
from mindspore import numpy as msnp

class Net(nn.Cell):
    def __init__(self, y, index, size):
        super(Net, self).__init__()
        self.y = y
        self.index = index
        self.norm = nn.Norm(-1)
        self.mask = msnp.zeros((y.shape[0],size))
        self.mask[msnp.arange(self.index.shape[0]),self.index] = 1

    def construct(self, x):
        return self.mask*self.y[self.index]*x

x = Tensor(np.array([1,2,3]).astype(np.float32))
y = Tensor(np.array([[1],[2],[3]]).astype(np.float32))
index = Tensor(np.array([0,2,1]).astype(np.int32))
shape = (y.shape[0], x.shape[0])

output = grad(Net(y,index,x.shape[0]))(ops.BroadcastTo(shape)(x))
print(output)
'''
[[1. 0. 0.]
 [0. 0. 3.]
 [0. 2. 0.]]
'''

这里我们看到得到的结果就是正确的了。当然，需要说明的是，虽然这个案例只是非常简单的内容，但是这里给出的如何去计算多维函数的自动微分的方法，同样也适用于一些更加复杂的网络和函数。

总结概要

在本文中通过一个实际函数案例的多次尝试，给出了得到预期结果的一种解决方案。虽然MindSpore框架本身提供了Jvp和Vjp等功能，但是实际上和Grad没有太大的区别，只是用Tuple的形式增加了输入的一个维度。如果可以使用纯Tensor的输入，用这种Mask加上Grad或者GradOperation的方案会更加简单一些。同时我也尝试过使用HyperMap（类似于Jax中的vmap）来解决这个问题，只需要写好一条对z求导的函数形式，就可以自动对这个求导过程进行扩维，两者的结果是一致的。但是MindSpore的HyperMap在Graph模式下兼容效果不是很好，建议非必要不尝试。