TensorFlow可微分编程实践3---向量微分和Jacobian矩阵

在这篇博文中，我们将利用TensorFlow Eager Execution API来实现一个完整多层感知器（MLP）模型。在具体实现多层感知器模型之前，我们首先来看，怎样用TensorFlow Eager Execution API来求向量与矩阵运算的导数。
我们知道在多层感知器模型中，最基本的运算是由第$l-1$层输出信号求出第$l$层神经元的输入信号，公式如下所示：

$$\begin{equation} \tag{3.2.001} \boldsymbol{z}^l=W^l\boldsymbol{a}^{l-1}+\boldsymbol{b}^l \end{equation}$$
为了下面讨论方便，我们假设第 $l-1$层有3个神经元，第 $l$层有2个神经元，式(3.2.001)中的各个值定义如下。
第 $l-1$层输出信号：
$$\begin{equation} \tag{3.2.002} \boldsymbol{a}^{l-1}=\begin{bmatrix} 1.0 \\ 2.0 \\ 3.0 \end{bmatrix} \end{equation}$$
第 $l-1$层到第 $l$层连接权值矩阵：
$$\begin{equation} \tag{3.2.003} W^l=\begin{bmatrix} 4.0 & 5.0 & 6.0 \\ 7.0 & 8.0 & 9.0 \end{bmatrix} \end{equation}$$
第 $l$层偏置值：
$$\begin{equation} \tag{3.2.004} \boldsymbol{b}^l=\begin{bmatrix} 1001.0 \\ 1002.0 \end{bmatrix} \end{equation}$$
为了进行学习，我们需要求出以下导数： $\frac{\partial{\boldsymbol{z}^l}}{\partial{\boldsymbol{a}^{l-1}}}$、 $\frac{\partial{\boldsymbol{z}^l}}{\partial{\boldsymbol{b}^{l}}}$、 $\frac{\partial{\boldsymbol{z}^l}}{\partial{W^l}}$，我们分别来进行讨论。
我们首先来看第一项，根据Jacobian矩阵定义得：
$$\begin{equation} \tag{3.2.005} \frac{\partial{\boldsymbol{z}^l}}{\partial{\boldsymbol{a}^{l-1}}}=\begin{bmatrix} \frac{\partial{z^l_1}}{\partial{a^{l-1}_1}} & \frac{\partial{z^l_1}}{\partial{a^{l-1}_2}} & ... & \frac{\partial{z^l_1}}{\partial{a^{l-1}_{N_{l-1}}}} \\ \frac{\partial{z^l_2}}{\partial{a^{l-1}_1}} & \frac{\partial{z^l_2}}{\partial{a^{l-1}_2}} & ... & \frac{\partial{z^l_2}}{\partial{a^{l-1}_{N_{l-1}}}} \\ ... & ... & ... & ... \\ \frac{\partial{z^l_{N_l}}}{\partial{a^{l-1}_1}} & \frac{\partial{z^l_{N_l}}}{\partial{a^{l-1}_2}} & ... & \frac{\partial{z^l_{N_l}}}{\partial{a^{l-1}_{N_{l-1}}}} \\ \end{bmatrix}=\begin{bmatrix} W^l_{1,1} & W^l_{1,2} & ... & W^l_{1,N_{l-1}} \\ W^l_{2,1} & W^l_{2,2} & ... & W^l_{2,N_{l-1}} \\ ... & ... & ... & ... \\ W^l_{N_l,1} & W^l_{N_l,2} & ... & W^l_{N_l,N_{l-1}} \\ \end{bmatrix}=W^l \end{equation}$$
我们接下来再来求对第 $l$层偏置求微分：
$$\begin{equation} \tag{3.2.006} \frac{\partial{\boldsymbol{z}^l}}{\boldsymbol{b}^l}=\begin{bmatrix} \frac{\partial{z^l_1}}{\partial{b^l_1}} & \frac{\partial{z^l_1}}{\partial{b^l_2}} & ... & \frac{\partial{z^l_1}}{\partial{b^l_{N_l}}} \\ \frac{\partial{z^l_2}}{\partial{b^l_1}} & \frac{\partial{z^l_2}}{\partial{b^l_2}} & ... & \frac{\partial{z^l_2}}{\partial{b^l_{N_l}}} \\ ... & ... & ... & ... \\ \frac{\partial{z^l_{N_l}}}{\partial{b^l_1}} & \frac{\partial{z^l_{N_l}}}{\partial{b^l_2}} & ... & \frac{\partial{z^l_{N_l}}}{\partial{b^l_{N_l}}} \\ \end{bmatrix}=\begin{bmatrix} 1 & 0 & ... & 0 \\ 0 & 1 & ... &０ \\ ... & ... & ... & ... \\ 0 & 0 & ... & 1 \end{bmatrix} \end{equation}$$
下面是一个向量对矩阵求偏导，而我们对这个操作没有定义，所以我们需要以一种变通的方式来进行，我们将 $W^l$视为由 $\boldsymbol{w}^{(i)} \in R^{N_{l-1}}$的行向量组成：
$$\begin{equation} \tag{3.2.007} W^l=\begin{bmatrix} (\boldsymbol{w}^{(1)})^T \\ (\boldsymbol{w}^{(2)})^T \\ ... \\ (\boldsymbol{w}^{(N_l)})^T \\ \end{bmatrix}=\boldsymbol{w}^l \in R^{N_l} \end{equation}$$
其实 $\boldsymbol{w}^{(i)}$是指向第 $l$层第 $i$个神经元所有连接权值组成的向量。
有了式(3.2.007)的定义，我们就可以将 $W^l$视为向量，这样根据Jacobian矩阵定义：
$$\begin{equation} \tag{3.2.008} \frac{\partial{\boldsymbol{z}^l}}{\partial{W^l}}=\frac{\partial{\boldsymbol{z}^l}}{\partial{\boldsymbol{w}^l}}=\begin{bmatrix} \frac{z^l_1}{\boldsymbol{w}^{(1)}} & \frac{z^l_1}{\boldsymbol{w}^{(2)}} & ... & \frac{z^l_1}{\boldsymbol{w}^{(N_l)}} \\ \frac{z^l_2}{\boldsymbol{w}^{(1)}} & \frac{z^l_2}{\boldsymbol{w}^{(2)}} & ... & \frac{z^l_2}{\boldsymbol{w}^{(N_l)}} \\ ... & ... & ... & ... \\ \frac{z^l_{N_l}}{\boldsymbol{w}^{(1)}} & \frac{z^l_{N_l}}{\boldsymbol{w}^{(2)}} & ... & \frac{z^l_{N_l}}{\boldsymbol{w}^{(N_l)}} \\ \end{bmatrix} \end{equation}$$
与前面不同的是，式(3.2.008)的矩阵中的每个元素都是一个标量对向量的求偏导，根据我们上篇博文介绍，标量对向量求偏导，结果为一个行向量，我们以 $\frac{z^l_i}{\boldsymbol{w}^{(j)}}$为例进行讨论。
如果 $i \neq j$时， $\boldsymbol{w}^{(j)}$是指向第 $l$层第 $j$个神经元的，不与第 $l$行第 $i$个神经元相接，因此所有偏层均为0，如下所示：
$$\begin{equation} \tag{3.2.009} \frac{z^l_i}{\boldsymbol{w}^{(j)}}=\begin{bmatrix} 0 & 0 & ... & 0 \end{bmatrix} \in R^{N_{l-1}} \end{equation}$$
如果 $i=j$时， $\boldsymbol{w}^{(j)}$是由指向第 $l$层第 $i$个神经元的所有连接权值组成的，根据输入信号定义可得：
$$\begin{equation} \tag{3.2.010} \frac{z^l_i}{\boldsymbol{w}^{(j)}}=\begin{bmatrix} \frac{\partial{z^l_i}}{\partial{W^l_{j,1}}} & \frac{\partial{z^l_i}}{\partial{W^l_{j,2}}} & ... & \frac{\partial{z^l_i}}{\partial{W^l_{j,N_{l-1}}}} \end{bmatrix}=\begin{bmatrix} a^{l-1}_1 & a^{l-1}_2 & ... & a^{l-1}_{N_{l-1}} \end{bmatrix} \end{equation}$$
因此式(3.2.008)矩阵的每个元素为指向纸里的 $R^{N_{l-1}}$向量，当不在对角线上时，所有元素值为零，当在对角线上时，元素为第 $l-1$层输出值。
下面我们来看，怎样通过TensorFlow Eager Excecution API来求出这些偏导的值。

@tf.custom_gradient
def f002(W, a, b):
    def grad_fn(dy):
        ws = W.shape
        pz_pW = np.zeros((2, 2, 3))
        a1 = tf.reshape(a, [3])
        for idx in range(ws[0]):
            pz_pW[idx][idx] = a1
        diag = tf.ones([W.shape[0]])
        d_b = tf.matrix_diag(diag)
        return tf.constant(pz_pW), W, d_b
    return tf.matmul(W, a) + b, grad_fn

def test001(args={}):
    tf.enable_eager_execution()
    tfe = tf.contrib.eager

    W = tf.constant([[4.0, 5.0, 6.0],[7.0, 8.0, 9.0]])
    a = tf.reshape(tf.constant([1.0, 2.0, 3.0]), [3, 1])
    i_debug = 2
    if 1 == i_debug:
        f003(W, a)
        return 
    b = tf.reshape(tf.constant([1001.0, 1002.0]), [2, 1])
    z = f002(W, a, b)
    print('z=Wa+b={0}'.format(z))
    grad_f1 = tfe.gradients_function(f002)
    dv = grad_f1(W, a, b)
    print('pz_pW={0}'.format(dv[0].numpy()))
    print('pz_pa={0}'.format(dv[1].numpy()))
    print('pz_pb={0}'.format(dv[2].numpy()))
    print('v0.0.1')

在上面的代码中，需要说明的是第16行，将向量$\boldsymbol{a}$定义为3行1列的矩阵形式，这是为了与连接权值矩阵做矩阵相乘。
在前向计算阶段，可以直接调用TensorFlow的矩阵乘法和加法，我们就可以取得正确的结果，但是在反向求导阶段，如果我们直接利用TensorFlow进行求导，例如求$\frac{\partial{\boldsymbol{z}^l}}{\partial{\boldsymbol{a}^{l-1}}}$时，TensorFlow返回结果维数与$\boldsymbol{a}^{l-1}$相同，而Jacobian矩阵的维数为$R^{N_l \times N_{l-1}}$，所以我们需要自己定义求导函数，根据上面的理论分析，$\frac{\partial{\boldsymbol{z}^l}}{\partial{W}^l}$是一个3维的张量，维数为$R^{N_l \times N_l \times N_{l-1}}$，我们可以将其视为一个$R^{N_l \times N_l}$的矩阵，矩阵中每个元素均为一个数组，长度为$N_{l-1}$，且除对角线上的元素外，数组元素为0，而对角线上的元素，数组元素为第$l-1$导神经元的输出值。$\frac{\partial{\boldsymbol{z}^l}}{\partial{\boldsymbol{a}^{l-1}}}$为第$l-1$层到第$l$层的连接权值矩阵$W^l$，而$\frac{\partial{\boldsymbol{z}^l}}{\partial{\boldsymbol{b}^{l}}}$为$R^{N_l \times N_l}$的单位阵，运行结果如下所示：

至此我们完成了第$l-1$层到第$l$层正向传输和反向求导工作，基本上按照数学理论要求，我们就可以完全处理一个多层感知器模型了。