LSTM背后的数学原理

引言

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

LSTM单元

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

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

遗忘门

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

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

公式为：

$${\mathbf{\Gamma }}_f^{⟨t⟩}=\sigma ({\mathbf{W}}_f[{\mathbf{a}}^{⟨t-1⟩},{\mathbf{x}}^{⟨t⟩}]+{\mathbf{b}}_f)\text{\hspace{1em}(1)}$$

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

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

候选值${\tilde{\mathbf{c}}}^{⟨t⟩}$

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

公式为：

$${\tilde{\mathbf{c}}}^{⟨t⟩}=\tanh \left({\mathbf{W}}_c[{\mathbf{a}}^{⟨t-1⟩},{\mathbf{x}}^{⟨t⟩}]+{\mathbf{b}}_c\right)\text{\hspace{1em}(2)}$$

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

cct代表${\tilde{\mathbf{c}}}^{⟨t⟩}$
Wc代表$W_c$

更新门(输入门)

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

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

公式：

$${\mathbf{\Gamma }}_i^{⟨t⟩}=\sigma ({\mathbf{W}}_i[a^{⟨t-1⟩},{\mathbf{x}}^{⟨t⟩}]+{\mathbf{b}}_i)\text{\hspace{1em}(3)}$$

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

单元状态$c^{⟨t⟩}$

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

公式：

$${\mathbf{c}}^{⟨t⟩}={\mathbf{\Gamma }}_f^{⟨t⟩}\ast {\mathbf{c}}^{⟨t-1⟩}+{\mathbf{\Gamma }}_i^{⟨t⟩}\ast {\tilde{\mathbf{c}}}^{⟨t⟩}\text{\hspace{1em}(4)}$$

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

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

c_next：当前单元状态$c^{⟨t⟩}$,形状$(n_a,m)$

c_prev: 前一个单元状态$c^{⟨t-1⟩}$,形状$(n_a,m)$

输出门${\mathbf{\Gamma }}_o$

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

公式：
$${\mathbf{\Gamma }}_o^{⟨t⟩}=\sigma ({\mathbf{W}}_o[{\mathbf{a}}^{⟨t-1⟩},{\mathbf{x}}^{⟨t⟩}]+{\mathbf{b}}_o)\text{\hspace{1em}(5)}$$

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

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

隐藏状态

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

公式：
$${\mathbf{a}}^{⟨t⟩}={\mathbf{\Gamma }}_o^{⟨t⟩}\ast \tanh ({\mathbf{c}}^{⟨t⟩})\text{\hspace{1em}(6)}$$

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

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

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

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

预测值${\hat{y}}^{⟨t⟩}$

• 在分类问题中的输出值使用softmax函数
  $$z^{⟨t⟩}={\mathbf{W}}_y{\mathbf{a}}^{⟨t⟩}+{\mathbf{b}}_y\text{\hspace{1em}(7)}$$

$${\mathbf{y}}_{pred}^{⟨t⟩}=\text{softmax}(z^{⟨t⟩})\text{\hspace{1em}(8)}$$

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

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

前向传播

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

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

$$concat=\left[\begin{array}{c}{a}^{⟨t-1⟩}\\ x^{⟨t⟩}\end{array}\right]$$

反向传播

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)\text{\hspace{1em}(9)}$$

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

$$d{z}^{⟨t⟩}={\hat{y}}^{⟨t⟩}-y^{⟨t⟩}\text{\hspace{1em}(10)}$$

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

$$d{W}_y=dz\cdot a^{⟨t⟩}\text{\hspace{1em}(11)}$$
$$d{b}_y=dz\text{\hspace{1em}(12)}$$

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

在时刻$T_x$时,
$$d{a}^{⟨T_x⟩}=\frac{\partial L}{\partial a^{⟨T_x⟩}}=\frac{\partial l(T_x)}{\partial a^{⟨T_x⟩}}=d{z}^{⟨T_x⟩}{W}_y\text{\hspace{1em}(13)}$$

$$d{c}^{⟨T_x⟩}=\frac{\partial L}{\partial c^{⟨T_x⟩}}=d{a}^{⟨T_x⟩}\cdot {\mathbf{\Gamma }}_o^{⟨T_x⟩}\cdot (1-tanh(c^{⟨T_x⟩}{)}^2)\text{\hspace{1em}(14)}$$

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

$$\begin{gathered} d{a}^{⟨t⟩}=\frac{\partial a^{⟨t+1⟩}}{a^{⟨t⟩}}+\frac{\partial l(t)}{\partial a^{⟨t⟩}} \\ =\frac{\partial L(t+1)}{\partial a^{⟨t+1⟩}}\frac{\partial a^{⟨t+1⟩}}{\partial a^{⟨t⟩}}+d{z}^{⟨t⟩}{W}_y \\ =d{z}^{⟨t+1⟩}{W}_y\cdot \frac{\partial a^{⟨t+1⟩}}{\partial a^{⟨t⟩}}+d{z}^{⟨t⟩}{W}_y\text{\hspace{1em}(15)} \end{gathered}$$

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

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

$$\begin{gathered} \frac{\partial a^{⟨t+1⟩}}{\partial a^{⟨t⟩}}=\frac{\partial a^{⟨t+1⟩}}{\partial {\mathbf{\Gamma }}_o^{⟨t+1⟩}}\frac{\partial {\mathbf{\Gamma }}_o^{⟨t+1⟩}}{\partial a^{⟨t⟩}}+\frac{\partial a^{⟨t+1⟩}}{\partial c^{⟨t+1⟩}}\frac{\partial c^{⟨t+1⟩}}{\partial {\tilde{\mathbf{c}}}^{⟨t+1⟩}}\frac{\partial {\tilde{\mathbf{c}}}^{⟨t+1⟩}}{\partial a^{⟨t⟩}}+\frac{\partial a^{⟨t+1⟩}}{\partial c^{⟨t+1⟩}}\frac{\partial c^{⟨t+1⟩}}{\partial {\mathbf{\Gamma }}_i^{⟨t+1⟩}}\frac{\partial {\mathbf{\Gamma }}_i^{⟨t+1⟩}}{\partial a^{⟨t⟩}}+\frac{\partial a^{⟨t+1⟩}}{\partial c^{⟨t+1⟩}}\frac{\partial c^{⟨t+1⟩}}{\partial {\mathbf{\Gamma }}_f^{⟨t+1⟩}}\frac{\partial {\mathbf{\Gamma }}_f^{⟨t+1⟩}}{\partial a^{⟨t⟩}} \\ =\tanh (c^{⟨t+1⟩})\cdot {\mathbf{\Gamma }}_o^{⟨t+1⟩}(1-{\mathbf{\Gamma }}_o^{⟨t+1⟩})W_o+{\mathbf{\Gamma }}_o^{⟨t+1⟩}(1-{\tanh (c^{⟨t+1⟩})}^2)\cdot {\mathbf{\Gamma }}_i^{⟨t+1⟩}\cdot (1-{{\tilde{\mathbf{c}}}^{⟨t⟩}}^2)W_c+{\mathbf{\Gamma }}_o^{⟨t+1⟩}(1-{\tanh (c^{⟨t+1⟩})}^2)\cdot {\tilde{\mathbf{c}}}^{⟨t+1⟩}\cdot {\mathbf{\Gamma }}_i^{⟨t⟩}(1-{\mathbf{\Gamma }}_i^{⟨t⟩})W_i+{\mathbf{\Gamma }}_o^{⟨t+1⟩}(1-{\tanh (c^{⟨t+1⟩})}^2)\cdot c^{⟨t⟩}\cdot {\mathbf{\Gamma }}_f^{⟨t⟩}(1-{\mathbf{\Gamma }}_f^{⟨t⟩})W_f\text{\hspace{1em}(16)} \end{gathered}$$

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

在时刻$t(t<T_x)$时，$c^{⟨t⟩}$的梯度也是由当前时刻的误差以及$t+1$时刻的误差组成(由公式$(4),(6)$)得：
$$\begin{gathered} d{c}^{⟨t⟩}=\frac{\partial L}{\partial c^{⟨t+1⟩}}\frac{\partial c^{⟨t+1⟩}}{\partial c^{⟨t⟩}}+\frac{\partial L}{\partial a^{⟨t⟩}}\frac{\partial a^{⟨t⟩}}{\partial c^{⟨t⟩}} \\ =d{c}^{⟨t+1⟩}\frac{\partial c^{⟨t+1⟩}}{\partial c^{⟨t⟩}}+d{a}^{⟨t⟩}{\mathbf{\Gamma }}_o^{⟨t⟩}(1-{\tanh (c^{⟨t⟩})}^2) \\ =d{c}^{⟨t+1⟩}{\mathbf{\Gamma }}_f^{⟨t+1⟩}+d{a}^{⟨t⟩}{\mathbf{\Gamma }}_o^{⟨t⟩}(1-{\tanh (c^{⟨t⟩})}^2)\text{\hspace{1em}(17)} \end{gathered}$$

---

现在求对$W_o,W_f,W_i,W_c$的梯度就简单了。
$$d{W}_o=\frac{\partial L}{\partial a^{⟨t⟩}}\cdot \frac{\partial a^{⟨t⟩}}{\partial {\mathbf{\Gamma }}_o^{⟨t⟩}}\cdot \frac{\partial {\mathbf{\Gamma }}_o^{⟨t⟩}}{W_o}=d{a}^{⟨t⟩}\cdot tanh(c^{⟨t⟩})\cdot {\Gamma }_o^{⟨t⟩}(1-{\Gamma }_o^{⟨t⟩}){\left[\begin{array}{c}{a}_{prev}\\ x_t\end{array}\right]}^T\text{\hspace{1em}(18)}$$
$$d{b}_o=d{a}^{⟨t⟩}\cdot \frac{\partial a^{⟨t⟩}}{\partial {\mathbf{\Gamma }}_o^{⟨t⟩}}\cdot \frac{\partial {\mathbf{\Gamma }}_o^{⟨t⟩}}{b_o}=d{a}^{⟨t⟩}\cdot tanh(c^{⟨t⟩})\cdot {\Gamma }_o^{⟨t⟩}(1-{\Gamma }_o^{⟨t⟩})\text{\hspace{1em}(19)}$$

$$\begin{gathered} d{W}_f=\frac{\partial L}{\partial c^{⟨t⟩}}\cdot \frac{\partial c^{⟨t⟩}}{\partial {\Gamma }_f^{⟨t⟩}}\cdot \frac{\partial {\Gamma }_f^{⟨t⟩}}{\partial W_f} \\ =d{c}^{⟨t⟩}\cdot c^{⟨t-1⟩}\cdot {\Gamma }_f^{⟨t⟩}(1-{\Gamma }_f^{⟨t⟩}){\left[\begin{array}{c}{a}_{prev}\\ x_t\end{array}\right]}^T\text{\hspace{1em}(20)} \end{gathered}$$

$$d{b}_f=d{c}^{⟨t⟩}\cdot c^{⟨t-1⟩}\cdot {\Gamma }_f^{⟨t⟩}(1-{\Gamma }_f^{⟨t⟩})\text{\hspace{1em}(21)}$$

$$\begin{gathered} d{W}_i=\frac{\partial L}{\partial c^{⟨t⟩}}\cdot \frac{\partial c^{⟨t⟩}}{\partial {\Gamma }_i^{⟨t⟩}}\cdot \frac{\partial {\Gamma }_i^{⟨t⟩}}{\partial W_i} \\ =d{c}^{⟨t⟩}\cdot {\tilde{\mathbf{c}}}^{⟨t⟩}\cdot {\Gamma }_i^{⟨t⟩}(1-{\Gamma }_i^{⟨t⟩}){\left[\begin{array}{c}{a}_{prev}\\ x_t\end{array}\right]}^T\text{\hspace{1em}(22)} \end{gathered}$$
$$d{b}_i=d{c}^{⟨t⟩}\cdot {\tilde{\mathbf{c}}}^{⟨t⟩}\cdot {\Gamma }_i^{⟨t⟩}(1-{\Gamma }_i^{⟨t⟩})\text{\hspace{1em}(23)}$$

$$\begin{gathered} d{W}_c=\frac{\partial L}{\partial c^{⟨t⟩}}\cdot \frac{\partial c^{⟨t⟩}}{\partial {\tilde{\mathbf{c}}}^{⟨t⟩}}\cdot \frac{\partial {\tilde{\mathbf{c}}}^{⟨t⟩}}{\partial W_c} \\ =d{c}^{⟨t⟩}\cdot {\Gamma }_i^{⟨t⟩}\cdot (1-{{\tilde{\mathbf{c}}}^{⟨t⟩}}^2){\left[\begin{array}{c}{a}_{prev}\\ x_t\end{array}\right]}^T\text{\hspace{1em}(24)} \end{gathered}$$

$$d{b}_c=d{c}^{⟨t⟩}\cdot {\Gamma }_i^{⟨t⟩}\cdot (1-{{\tilde{\mathbf{c}}}^{⟨t⟩}}^2)\text{\hspace{1em}(25)}$$

参考

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