LSTM背后的数学原理

引言

LSTM是RNN的变种，是为了解决RNN存在的长期依赖问题而专门设计出来的。所谓长期依赖问题是，后面的单词在很长的时间序列后还依赖前面的单词，但由于梯度消失问题，导致前面的单词无法影响到后面的单词。

LSTM单元

LSTM单元(cell)在每个时间点更新单元状态$c^{\langle t \rangle}$，它决定了$a^{\langle t \rangle}$的值。LSTM有更新门、遗忘门和输出门来控制这些值。

下面来对LSTM中的元素做一些说明

遗忘门

遗忘门用来控制内存中之前的状态是否会被遗忘掉。

• 如果遗忘门的值是0，LSTM会遗忘(忽略)之前的状态
• 如果遗忘门的值是1，LSTM会记得(保持)之前的状态
• 如果是0到1之间的值，代表LSTM会记得之前的状态多大程度

公式为：

$\mathbf{\Gamma}_f^{\langle t \rangle} = \sigma(\mathbf{W}_f[\mathbf{a}^{\langle t-1 \rangle}, \mathbf{x}^{\langle t \rangle}] + \mathbf{b}_f)\tag{1}$

• $W_f$和$b_f$是可学习的权重和偏差
• 通过sigmoid函数来保证输出的值在[0,1]之间
• 遗忘门$\mathbf{\Gamma}_f^{\langle t \rangle}$与之前单元状态$c^{\langle t \rangle}$同维度，即它们能逐元素相乘

在代码中​​Wf​​​代表$W_f$,​​​bf​​​代表$b_f$,​​​ft​​​代表$\mathbf{\Gamma}_f^{\langle t \rangle}$

候选值 c ~ ⟨ t ⟩ \tilde{\mathbf{c}}^{\langle t \rangle} c~⟨t⟩

• 候选值保存的是当前时间点可能会存入当前单元状态($c^{\langle t \rangle}$)的信息
• 候选值能多大程度的存入当前单元状态取决于更新门

公式为：

$\mathbf{\tilde{c}}^{\langle t \rangle} = \tanh\left( \mathbf{W}_{c} [\mathbf{a}^{\langle t - 1 \rangle}, \mathbf{x}^{\langle t \rangle}] + \mathbf{b}_{c} \right) \tag{2}$

这里用的是tanh函数，所以取值范围为[-1,1]

​​cct​​​代表$\tilde{\mathbf{c}}^{\langle t \rangle}$
​​​Wc​​​代表$W_c$

更新门(输入门)

• 更新门决定候选值(哪些维度)能多大程度的存入当前单元状态
• 如果更新门的值是0，意味着防止候选值存入单元状态
• 如果更新门的值是1，意味着完全允许候选值存入单元状态

有些文献称它为输入门，并且用"i"来表示，这里沿用这种约定

公式：

$\mathbf{\Gamma}_i^{\langle t \rangle} = \sigma(\mathbf{W}_i[a^{\langle t-1 \rangle}, \mathbf{x}^{\langle t \rangle}] + \mathbf{b}_i)\tag{3}$

​​Wi​​​代表$W_i$,​​​bi​​​代表$b_i$,​​​it​​​代表更新门$\mathbf{\Gamma}_i^{\langle t \rangle}$。

单元状态 c ⟨ t ⟩ c^{\langle t \rangle} c⟨t⟩

• 单元状态是时间序列间传递的"记忆"
• 新的单元状态由之前的状态和当前候选值组成

公式：

$\mathbf{c}^{\langle t \rangle} = \mathbf{\Gamma}_f^{\langle t \rangle}* \mathbf{c}^{\langle t-1 \rangle} + \mathbf{\Gamma}_{i}^{\langle t \rangle} *\mathbf{\tilde{c}}^{\langle t \rangle} \tag{4}$

• 结合上面所有的公式，得到了单元状态的计算公式
• 前一单元状态由遗忘门控制会有多少被保存到当前单元状态中
• 候选值由更新门控制能有多少被保存到当前单元状态中

​​c​​​：所有时间点的单元状态$c$，形状是$(n_a,m,T)$

​​c_next​​​：当前单元状态$c^{\langle t \rangle}$,形状$(n_a,m)$

​​c_prev​​​: 前一个单元状态$c^{\langle t-1 \rangle}$,形状$(n_a,m)$

输出门 Γ o \mathbf{\Gamma}_{o} Γo​

• 输出门控制了当前时间点能输出什么
• 和之前所有门一样，取值范围[0,1]

公式：
$\mathbf{\Gamma}_o^{\langle t \rangle}= \sigma(\mathbf{W}_o[\mathbf{a}^{\langle t-1 \rangle}, \mathbf{x}^{\langle t \rangle}] + \mathbf{b}_{o})\tag{5}$

​​W_o​​​代表输出门的权重$W_o$,​​​bo​​​代表输出门的偏差$b_o$,​​​ot​​​代表输出门$\mathbf{\Gamma}_{o}$

从三个门的公式可以看出，它们的激活函数都是sigmoid，取值都是[0,1]，输入都是$a^{\langle t-1 \rangle}$和$x^{\langle t \rangle}$，唯一的区别是可学习的权重和偏差不一样。如果取值为0，表示这个门是关闭的；取值为1，表示这个门是完全打开的；取值$(0,1)$表示这个门是半关半开的，只允许一部分的值进入(被保存,被传递)。

隐藏状态

• 当前的隐藏状态会传递到下一个时间点的LSTM单元
• 它用于决定下个时间点的三个门
• 同时也用于当前时间点的预测(输出值$\hat y^{\langle t \rangle}$)

公式：
$\mathbf{a}^{\langle t \rangle} = \mathbf{\Gamma}_o^{\langle t \rangle} * \tanh(\mathbf{c}^{\langle t \rangle})\tag{6}$

• 隐藏状态由单元状态和输出门决定
• 单元状态传递到tanh函数得到$[-1,1]$的取值

​​a​​​： 所有的隐藏状态$a$，形状$(n_a,m,T_x)$

​​a_prev​​​: 上个时间点的隐藏状态$a^{\langle t-1 \rangle}$，形状$(n_a,m)$

​​a_next​​​: 当前时间点的隐藏状态$a^{\langle t \rangle}$，形状$(n_a,m)$

预测值 y ^ ⟨ t ⟩ \hat y^{\langle t \rangle} y^​⟨t⟩

• 在分类问题中的输出值使用softmax函数
  $z^{\langle t \rangle} = \mathbf{W}_{y} \mathbf{a}^{\langle t \rangle} + \mathbf{b}_{y} \tag{7}$

$\mathbf{y}^{\langle t \rangle}_{pred} = \textrm{softmax}(z^{\langle t \rangle}) \tag{8}$

​​y_pred​​​: 所有时间点的预测值$y_{pred}$，形状$(n_y,m,T_x)$
​​​yt_pred​​​: 当前时间点的预测值$y_{pred}^{\langle t \rangle}$,形状$(n_y,m)$

至此我们知道了LSTM单元中的所有计算公式，下面来看如何实现前向传播和反向传播。

前向传播

实现如上图所示的前向传播过程，我们需要代码化上面的公式$(1)$~$(7)$。

要注意的是，我们会叠加前一个隐藏状态$a^{\langle t-1 \rangle}$和当前的输入$x^{\langle t \rangle}$到一个矩阵​​​concat​​：

$concat = \begin{bmatrix} a^{\langle t-1 \rangle} \\ x^{\langle t \rangle} \end{bmatrix}$

反向传播

LSTM的反向传播比RNN的要复杂一点。不过遵循规则——求某个节点的梯度时，考虑该节点的所有输出节点。分别计算每个输出节点的梯度乘上输出节点对该节点的梯度，然后加起来就得到该节点的梯度，也不难。

首先列出激活函数的导数：
$d \tanh(x) = 1 - \tanh(x)^2$
$d \sigma(x) = \sigma(x)(1 - \sigma(x))$

假设考虑的LSTM结构为多对多的，且$T_x=T_y$，每个时刻$t$都有一个输出及一个损失$l(t)$，全局损失函数为：
$L = \sum_{i=1}^{T_x} l(t) \tag{9}$

我们求$L$对$z^{\langle t \rangle}$的导数$dz^{\langle t \rangle}$，具体过程可以参考博客 ​​​Softmax与Cross-entropy的求导​​，得到：

$dz^{\langle t \rangle} = \hat y^{\langle t \rangle} - y^{\langle t \rangle} \tag{10}$

根据公式$(7)$，可以很容易的求出：

$dW_y = dz \cdot a^{\langle t \rangle} \tag{11}$
$db_y = dz \tag{12}$

而求$da$和$dc$时要分两种情况考虑：

在时刻$T_x$时,
$da^{\langle T_x \rangle} = \frac{\partial L}{\partial a^{\langle T_x \rangle}}= \frac{\partial l(T_x)}{\partial a^{\langle T_x \rangle}} = dz^{\langle T_x \rangle} W_y \tag{13}$

$dc^{\langle T_x \rangle} = \frac{\partial L}{\partial c^{\langle T_x \rangle}} = da^{\langle T_x \rangle}\cdot \mathbf{\Gamma}_o^{\langle T_x \rangle} \cdot (1 - tanh(c^{\langle T_x \rangle})^2) \tag{14}$

在时刻$t \,\,(t < T_x)$时，$a^{\langle t \rangle}$的后续同时有$a^{\langle t+1 \rangle}$(大于$t$时刻的误差)和$y^{\langle t \rangle}$($t$时刻的误差)两个节点。因此计算梯度时要考虑这两部分：

$da^{\langle t \rangle} = \frac{\partial a^{\langle t+1 \rangle}}{a^{\langle t \rangle}} + \frac{\partial l(t)}{\partial a^{\langle t \rangle}} \\ = \frac{\partial L(t+1)}{\partial a^{\langle t+1 \rangle}}\frac{\partial a^{\langle t+1 \rangle}}{\partial a^{\langle t \rangle}} + dz^{\langle t \rangle} W_y \\ = dz^{\langle t+1 \rangle} W_y \cdot \frac{\partial a^{\langle t+1 \rangle}}{\partial a^{\langle t \rangle}} +dz^{\langle t \rangle} W_y \tag{15}$

在这一步反向传播计算的难点在于$\frac{\partial a^{\langle t+1 \rangle}}{\partial a^{\langle t \rangle}}$。

因为$a^{\langle t \rangle}$受到上图这四部分所影响，而这四部分都和$a^{\langle t-1 \rangle}$有关。所以$\frac{\partial a^{\langle t+1 \rangle}}{\partial a^{\langle t \rangle}}$的计算结果也由四部分组成(公式$(6,5),(6,4,2),(6,4,3),(6,4,1)$)：

$\frac{\partial a^{\langle t+1 \rangle}}{\partial a^{\langle t \rangle}} = \frac{\partial a^{\langle t+1 \rangle}}{\partial \mathbf{\Gamma}_o^{\langle t+1 \rangle}} \frac{\partial \mathbf{\Gamma}_o^{\langle t+1 \rangle} }{\partial a^{\langle t \rangle} } + \frac{\partial a^{\langle t+1 \rangle}}{\partial c^{\langle t+1 \rangle}} \frac{\partial c^{\langle t+1 \rangle} }{\partial \mathbf{\tilde{c}}^{\langle t+1 \rangle} } \frac{\partial \mathbf{\tilde{c}}^{\langle t+1 \rangle} }{\partial a^{\langle t \rangle}} + \frac{\partial a^{\langle t+1 \rangle}}{\partial c^{\langle t+1 \rangle}} \frac{\partial c^{\langle t+1 \rangle} }{\partial \mathbf{\Gamma}_i^{\langle t+1 \rangle}} \frac{\partial \mathbf{\Gamma}_i^{\langle t+1 \rangle} }{\partial a^{\langle t \rangle}} + \frac{\partial a^{\langle t+1 \rangle}}{\partial c^{\langle t+1 \rangle}} \frac{\partial c^{\langle t+1 \rangle} }{\partial \mathbf{\Gamma}_f^{\langle t+1 \rangle}} \frac{\partial \mathbf{\Gamma}_f^{\langle t+1 \rangle} }{\partial a^{\langle t \rangle}} \\ = \tanh(c^{\langle t+1 \rangle}) \cdot \mathbf{\Gamma}_o^{\langle t+1 \rangle}(1-\mathbf{\Gamma}_o^{\langle t+1 \rangle})W_o + \mathbf{\Gamma}_o^{\langle t+1 \rangle}(1-{\tanh(c^{\langle t+1 \rangle})}^2)\cdot \mathbf{\Gamma}_i^{\langle t+1 \rangle}\cdot (1-{\mathbf{\tilde{c}}^{\langle t \rangle}}^2)W_c + \mathbf{\Gamma}_o^{\langle t+1 \rangle}(1-{\tanh(c^{\langle t+1 \rangle})}^2) \cdot \mathbf{\tilde{c}}^{\langle t+1 \rangle} \cdot \mathbf{\Gamma}_i^{\langle t \rangle}(1 - \mathbf{\Gamma}_i^{\langle t \rangle})W_i + \mathbf{\Gamma}_o^{\langle t+1 \rangle}(1-{\tanh(c^{\langle t+1 \rangle})}^2) \cdot c^{\langle t \rangle} \cdot \mathbf{\Gamma}_f^{\langle t \rangle}(1-\mathbf{\Gamma}_f^{\langle t \rangle})W_f \tag{16}$

上面有一个公共项$\mathbf{\Gamma}_o^{\langle t+1 \rangle}(1-{\tanh(c^{\langle t+1 \rangle})}^2)$

在时刻$t \,\,(t < T_x)$时，$c^{\langle t \rangle}$的梯度也是由当前时刻的误差以及$t+1$时刻的误差组成(由公式$(4),(6)$)得：
$dc^{\langle t \rangle} = \frac{\partial L}{\partial c^{\langle t+1 \rangle}} \frac{\partial c^{\langle t+1 \rangle}}{\partial c^{\langle t \rangle}} + \frac{\partial L}{\partial a^{\langle t \rangle}} \frac{\partial a^{\langle t \rangle}}{\partial c^{\langle t \rangle}} \\ = dc^{\langle t+1 \rangle}\frac{\partial c^{\langle t+1 \rangle}}{\partial c^{\langle t \rangle}} + da^{\langle t \rangle}\mathbf{\Gamma}_o^{\langle t \rangle}(1-{\tanh(c^{\langle t \rangle})}^2) \\ = dc^{\langle t+1 \rangle}\mathbf{\Gamma}_f^{\langle t+1 \rangle} + da^{\langle t \rangle}\mathbf{\Gamma}_o^{\langle t \rangle}(1-{\tanh(c^{\langle t \rangle})}^2) \tag{17}$

---

现在求对$W_o,W_f,W_i,W_c$的梯度就简单了。
$dW_o =\frac{\partial L}{\partial a^{\langle t \rangle}} \cdot \frac{\partial a^{\langle t \rangle}}{\partial \mathbf{\Gamma}_o^{\langle t \rangle}} \cdot \frac{\partial \mathbf{\Gamma}_o^{\langle t \rangle}}{W_o} = da^{\langle t \rangle} \cdot tanh(c^{\langle t \rangle}) \cdot \Gamma_o^{\langle t \rangle}(1-\Gamma_o^{\langle t \rangle}) \begin{bmatrix} a_{prev} \\ x_t\end{bmatrix}^T \tag{18}$
$db_o = da^{\langle t \rangle} \cdot \frac{\partial a^{\langle t \rangle}}{\partial \mathbf{\Gamma}_o^{\langle t \rangle}} \cdot \frac{\partial \mathbf{\Gamma}_o^{\langle t \rangle}}{b_o} = da^{\langle t \rangle} \cdot tanh(c^{\langle t \rangle}) \cdot \Gamma_o^{\langle t \rangle}(1-\Gamma_o^{\langle t \rangle}) \tag{19}$

$dW_f = \frac{\partial L}{\partial c^{\langle t \rangle}} \cdot \frac{\partial c^{\langle t \rangle}}{\partial \Gamma_f^{\langle t \rangle}} \cdot \frac{\partial \Gamma_f^{\langle t \rangle}}{\partial W_f} \\ = dc^{\langle t \rangle} \cdot c^{\langle t-1 \rangle} \cdot \Gamma_f^{\langle t \rangle}(1-\Gamma_f^{\langle t \rangle})\begin{bmatrix} a_{prev} \\ x_t\end{bmatrix}^T \tag{20}$

$db_f = dc^{\langle t \rangle} \cdot c^{\langle t-1 \rangle} \cdot \Gamma_f^{\langle t \rangle}(1-\Gamma_f^{\langle t \rangle}) \tag{21}$

$dW_i = \frac{\partial L}{\partial c^{\langle t \rangle}} \cdot \frac{\partial c^{\langle t \rangle}}{\partial \Gamma_i^{\langle t \rangle}} \cdot \frac{\partial \Gamma_i^{\langle t \rangle}}{\partial W_i} \\ = dc^{\langle t \rangle} \cdot \mathbf{\tilde{c}}^{\langle t \rangle} \cdot \Gamma_i^{\langle t \rangle}(1-\Gamma_i^{\langle t \rangle})\begin{bmatrix} a_{prev} \\ x_t\end{bmatrix}^T \tag{22}$
$db_i = dc^{\langle t \rangle} \cdot \mathbf{\tilde{c}}^{\langle t \rangle} \cdot \Gamma_i^{\langle t \rangle}(1-\Gamma_i^{\langle t \rangle}) \tag{23}$

$dW_c = \frac{\partial L}{\partial c^{\langle t \rangle}} \cdot \frac{\partial c^{\langle t \rangle}}{\partial \mathbf{\tilde{c}}^{\langle t \rangle}} \cdot \frac{\partial \mathbf{\tilde{c}}^{\langle t \rangle}}{\partial W_c} \\ = dc^{\langle t \rangle} \cdot \Gamma_i^{\langle t \rangle}\cdot (1-{\mathbf{\tilde{c}}^{\langle t \rangle}}^2) \begin{bmatrix} a_{prev} \\ x_t\end{bmatrix}^T \tag{24}$

$db_c = dc^{\langle t \rangle} \cdot \Gamma_i^{\langle t \rangle}\cdot (1-{\mathbf{\tilde{c}}^{\langle t \rangle}}^2) \tag{25}$

参考

1. ​​LSTM模型与前向反向传播算法​​
2. ​​从零实现循环神经网路​​
3. 吴恩达课程