Tensorflow2.0之LSTM的原理与实现

LSTM 原理

LSTM 中引入了3个门，即输入门（input gate）、遗忘门（forget gate）和输出门（output gate），以及与隐藏状态形状相同的记忆细胞（某些文献把记忆细胞当成一种特殊的隐藏状态），从而记录额外的信息。

输入门、遗忘门和输出门

与门控循环单元中的重置门和更新门一样，如下图所示，长短期记忆的门的输入均为当前时间步输入${\mathbf{X}}_t$ 与上一时间步隐藏状态 ${\mathbf{H}}_{t-1}$，输出由激活函数为 sigmoid 函数的全连接层计算得到。如此一来，这3个门元素的值域均为$[0,1]$。

具体来说，假设样本数为 $n$，输入个数为 $d$，隐藏单元个数为 $h$，给定时间步 $t$ 的小批量输入 ${\mathbf{X}}_t\in {\mathbb{R}}^{n\times d}$ 和上一时间步隐藏状态 ${\mathbf{H}}_{t-1}\in {\mathbb{R}}^{n\times h}$。 时间步 $t$ 的输入门 ${\mathbf{I}}_t\in {\mathbb{R}}^{n\times h}$、遗忘门 ${\mathbf{F}}_t\in {\mathbb{R}}^{n\times h}$ 和输出门 ${\mathbf{O}}_t\in {\mathbb{R}}^{n\times h}$ 分别计算如下：

$${\mathbf{I}}_t=\sigma ({\mathbf{X}}_t{\mathbf{W}}_{xi}+{\mathbf{H}}_{t-1}{\mathbf{W}}_{hi}+{\mathbf{b}}_i)$$
$${\mathbf{F}}_t=\sigma ({\mathbf{X}}_t{\mathbf{W}}_{xf}+{\mathbf{H}}_{t-1}{\mathbf{W}}_{hf}+{\mathbf{b}}_f)$$
$${\mathbf{O}}_t=\sigma ({\mathbf{X}}_t{\mathbf{W}}_{xo}+{\mathbf{H}}_{t-1}{\mathbf{W}}_{ho}+{\mathbf{b}}_o)$$

其中的 ${\mathbf{W}}_{xi},{\mathbf{W}}_{xf},{\mathbf{W}}_{xo}\in {\mathbb{R}}^{d\times h}$ 和 ${\mathbf{W}}_{hi},{\mathbf{W}}_{hf},{\mathbf{W}}_{ho}\in {\mathbb{R}}^{h\times h}$ 是权重参数，${\mathbf{b}}_i,{\mathbf{b}}_f,{\mathbf{b}}_o\in {\mathbb{R}}^{1\times h}$ 是偏差参数。

候选记忆细胞

接下来，长短期记忆需要计算候选记忆细胞 ${\tilde{\mathbf{C}}}_t$。它的计算与上面介绍的3个门类似，但使用了值域在$[-1,1]$的 tanh 函数作为激活函数，如下图所示。

具体来说，时间步 $t$ 的候选记忆细胞 ${\tilde{\mathbf{C}}}_t\in {\mathbb{R}}^{n\times h}$ 的计算为

$${\tilde{\mathbf{C}}}_t=\text{tanh}({\mathbf{X}}_t{\mathbf{W}}_{xc}+{\mathbf{H}}_{t-1}{\mathbf{W}}_{hc}+{\mathbf{b}}_c)$$

其中 ${\mathbf{W}}_{xc}\in {\mathbb{R}}^{d\times h}$ 和 ${\mathbf{W}}_{hc}\in {\mathbb{R}}^{h\times h}$ 是权重参数，${\mathbf{b}}_c\in {\mathbb{R}}^{1\times h}$ 是偏差参数。

记忆细胞

我们可以通过元素值域在$[0,1]$的输入门、遗忘门和输出门来控制隐藏状态中信息的流动，这一般也是通过使用按元素乘法（符号为$\odot$）来实现的。当前时间步记忆细胞 ${\mathbf{C}}_t\in {\mathbb{R}}^{n\times h}$ 的计算组合了上一时间步记忆细胞和当前时间步候选记忆细胞的信息，并通过遗忘门和输入门来控制信息的流动：

$${\mathbf{C}}_t={\mathbf{F}}_t\odot {\mathbf{C}}_{t-1}+{\mathbf{I}}_t\odot {\tilde{\mathbf{C}}}_t.$$

如下图所示，遗忘门控制上一时间步的记忆细胞 ${\mathbf{C}}_{t-1}$ 中的信息是否传递到当前时间步，而输入门则控制当前时间步的输入 ${\mathbf{X}}_t$ 通过候选记忆细胞 ${\tilde{\mathbf{C}}}_t$ 如何流入当前时间步的记忆细胞。如果遗忘门一直近似1且输入门一直近似0，过去的记忆细胞将一直通过时间保存并传递至当前时间步。这个设计可以应对循环神经网络中的梯度衰减问题，并更好地捕捉时间序列中时间步距离较大的依赖关系。

隐藏状态

有了记忆细胞以后，接下来我们还可以通过输出门来控制从记忆细胞到隐藏状态 ${\mathbf{H}}_t\in {\mathbb{R}}^{n\times h}$ 的信息的流动：

$${\mathbf{H}}_t={\mathbf{O}}_t\odot \text{tanh}({\mathbf{C}}_t).$$

这里的 tanh 函数确保隐藏状态元素值在-1到1之间。需要注意的是，当输出门近似1时，记忆细胞信息将传递到隐藏状态供输出层使用；当输出门近似0时，记忆细胞信息只自己保留。下图展示了长短期记忆中隐藏状态的计算。

输出结果

长短期记忆的隐藏层输出包括隐藏状态和记忆细胞，但只有隐藏状态会传递到输出层，而记忆细胞不参与输出层的计算。
将上面得到的隐藏状态 $H_t$ 输入神经网络，得到结果 ${\mathbf{Y}}_t\in {\mathbb{R}}^{n\times q}$：
$${\mathbf{Y}}_t={\mathbf{H}}_t{\mathbf{W}}_{hq}+{\mathbf{b}}_q$$

低级API实现LSTM

在Tensorflow2.0之从零开始实现循环神经网络中，我们已经实现了用低级 API 实现 RNN，在这里，我们只需要将其中的三个函数稍作修改即可。

1、修改定义参数函数

将 get_params() 函数修改为：

def get_params():
    def _one(shape):
        return tf.Variable(tf.random.normal(shape=shape,stddev=0.01,mean=0,dtype=tf.float32))
    def _three():
        return (_one((num_inputs, num_hiddens)),
                _one((num_hiddens, num_hiddens)),
                tf.Variable(tf.zeros(num_hiddens), dtype=tf.float32))
    W_xi, W_hi, b_i = _three()  # 输入门参数
    W_xf, W_hf, b_f = _three()  # 遗忘门参数
    W_xo, W_ho, b_o = _three()  # 输出门参数
    W_xc, W_hc, b_c = _three()  # 候选记忆细胞参数
    
    # 输出层参数
    W_hq = _one((num_hiddens, num_outputs))
    b_q = tf.Variable(tf.zeros(num_outputs), dtype=tf.float32)
    return [W_xi, W_hi, b_i, W_xf, W_hf, b_f, W_xo, W_ho, b_o, W_xc, W_hc, b_c, W_hq, b_q]

2、修改初始化函数

在初始化函数中，长短期记忆的隐藏状态需要返回额外的形状为 (批量大小, 隐藏单元个数) 的值为0的记忆细胞。

def init_lstm_state(batch_size, num_hiddens):
    return (tf.zeros(shape=(batch_size, num_hiddens)), 
            tf.zeros(shape=(batch_size, num_hiddens)))

3、修改定义模型函数

def lstm(inputs, state, params):
    W_xi, W_hi, b_i, W_xf, W_hf, b_f, W_xo, W_ho, b_o, W_xc, W_hc, b_c, W_hq, b_q = params
    (H, C) = state
    outputs = []
    for X in inputs:
        X=tf.reshape(X,[-1,W_xi.shape[0]])
        I = tf.sigmoid(tf.matmul(X, W_xi) + tf.matmul(H, W_hi) + b_i)
        F = tf.sigmoid(tf.matmul(X, W_xf) + tf.matmul(H, W_hf) + b_f)
        O = tf.sigmoid(tf.matmul(X, W_xo) + tf.matmul(H, W_ho) + b_o)
        C_tilda = tf.tanh(tf.matmul(X, W_xc) + tf.matmul(H, W_hc) + b_c)
        C = F * C + I * C_tilda
        H = O * tf.tanh(C)
        Y = tf.matmul(H, W_hq) + b_q
        outputs.append(Y)
    return outputs, (H, C)

Tensorflow2.0 实现LSTM

在Tensorflow2.0之用循环神经网络生成周杰伦歌词中，我们已经用 Tensorflow2.0 中提供的高级 API 实现了循环神经网络，在这里，我们只需要修改实例化 RNN 的部分即可。
也就是将：

cell = keras.layers.SimpleRNNCell(num_hiddens, 
                                  kernel_initializer='glorot_uniform')
rnn_layer = keras.layers.RNN(cell,time_major=True,
                             return_sequences=True,return_state=True)

修改为：

rnn_layer = keras.layers.LSTM(num_hiddens,time_major=True,
                              return_sequences=True,return_state=True)