Run this notebook online:Binder or Colab: Colab

9.2. 长短期记忆网络（LSTM）

长期以来，隐变量模型存在着长期信息保存和短期输入缺失的问题。 解决这一问题的最早方法之一是长短期存储器（long short-term memory，LSTM） [Hochreiter & Schmidhuber, 1997]。 它有许多与门控循环单元（ Section 9.1）一样的属性。 有趣的是，长短期记忆网络的设计比门控循环单元稍微复杂一些， 却比门控循环单元早诞生了近20年。

9.2.1. 门控记忆元

可以说，长短期记忆网络的设计灵感来自于计算机的逻辑门。 长短期记忆网络引入了记忆元（memory cell），或简称为单元（cell）。 有些文献认为记忆元是隐状态的一种特殊类型， 它们与隐状态具有相同的形状，其设计目的是用于记录附加的信息。 为了控制记忆元，我们需要许多门。 其中一个门用来从单元中输出条目，我们将其称为输出门（output gate）。 另外一个门用来决定何时将数据读入单元，我们将其称为输入门（input gate）。 我们还需要一种机制来重置单元的内容，由遗忘门（forget gate）来管理， 这种设计的动机与门控循环单元相同， 能够通过专用机制决定什么时候记忆或忽略隐状态中的输入。 让我们看看这在实践中是如何运作的。

9.2.1.1. 输入门、忘记门和输出门

就如在门控循环单元中一样， 当前时间步的输入和前一个时间步的隐状态 作为数据送入长短期记忆网络的门中， 如 lstm_0所示。 它们由三个具有sigmoid激活函数的全连接层处理， 以计算输入门、遗忘门和输出门的值。 因此，这三个门的值都在\((0, 1)\)的范围内。

长短期记忆模型中的输入门、遗忘门和输出门 .. _lstm_0:

我们来细化一下长短期记忆网络的数学表达。 假设有\(h\)个隐藏单元，批量大小为\(n\)，输入数为\(d\)。 因此，输入为\(\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}\)。 它们的计算方法如下：

(9.2.1)\[\begin{split}\begin{aligned} \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), \end{aligned}\end{split}\]

其中\(\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}\)是偏置参数。

9.2.1.2. 候选记忆元

由于还没有指定各种门的操作，所以先介绍候选记忆元（candidate memory cell） \(\tilde{\mathbf{C}}_t \in \mathbb{R}^{n \times h}\)。 它的计算与上面描述的三个门的计算类似， 但是使用\(\tanh\)函数作为激活函数，函数的值范围为\((-1, 1)\)。 下面导出在时间步\(t\)处的方程：

(9.2.2)\[\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}\)是偏置参数。

候选记忆元的如 lstm_1所示。

长短期记忆模型中的候选记忆元 .. _lstm_1:

9.2.1.3. 记忆元

在门控循环单元中，有一种机制来控制输入和遗忘（或跳过）。 类似地，在长短期记忆网络中，也有两个门用于这样的目的： 输入门\(\mathbf{I}_t\)控制采用多少来自\(\tilde{\mathbf{C}}_t\)的新数据， 而遗忘门\(\mathbf{F}_t\)控制保留多少过去的 记忆元\(\mathbf{C}_{t-1} \in \mathbb{R}^{n \times h}\)的内容。 使用按元素乘法，得出：

(9.2.3)\[\mathbf{C}_t = \mathbf{F}_t \odot \mathbf{C}_{t-1} + \mathbf{I}_t \odot \tilde{\mathbf{C}}_t.\]

如果遗忘门始终为\(1\)且输入门始终为\(0\)， 则过去的记忆元\(\mathbf{C}_{t-1}\) 将随时间被保存并传递到当前时间步。 引入这种设计是为了缓解梯度消失问题， 并更好地捕获序列中的长距离依赖关系。

这样我们就得到了计算记忆元的流程图，如 Fig. 9.2.1。

https://d2l.ai/_images/lstm-2.svg

Fig. 9.2.1 在长短期记忆网络模型中计算记忆元

9.2.1.4. 隐状态

最后，我们需要定义如何计算隐状态 \(\mathbf{H}_t \in \mathbb{R}^{n \times h}\)， 这就是输出门发挥作用的地方。 在长短期记忆网络中，它仅仅是记忆元的\(\tanh\)的门控版本。 这就确保了\(\mathbf{H}_t\)的值始终在区间\((-1, 1)\)内：

(9.2.4)\[\mathbf{H}_t = \mathbf{O}_t \odot \tanh(\mathbf{C}_t).\]

只要输出门接近\(1\)，我们就能够有效地将所有记忆信息传递给预测部分， 而对于输出门接近\(0\)，我们只保留记忆元内的所有信息，而不需要更新隐状态。

lstm_3提供了数据流的图形化演示。

在长短期记忆模型中计算隐状态 .. _lstm_3:

9.2.2. 从零开始实现

现在，我们从零开始实现长短期记忆网络。 与 Section 8.5中的实验相同， 我们首先加载时光机器数据集。

%load ../utils/djl-imports
%load ../utils/plot-utils
%load ../utils/Functions.java
%load ../utils/PlotUtils.java

%load ../utils/StopWatch.java
%load ../utils/Accumulator.java
%load ../utils/Animator.java
%load ../utils/Training.java
%load ../utils/timemachine/Vocab.java
%load ../utils/timemachine/RNNModel.java
%load ../utils/timemachine/RNNModelScratch.java
%load ../utils/timemachine/TimeMachine.java
%load ../utils/timemachine/TimeMachineDataset.java

NDManager manager = NDManager.newBaseManager();

int batchSize = 32;
int numSteps = 35;

TimeMachineDataset dataset =
        new TimeMachineDataset.Builder()
                .setManager(manager)
                .setMaxTokens(10000)
                .setSampling(batchSize, false)
                .setSteps(numSteps)
                .build();
dataset.prepare();
Vocab vocab = dataset.getVocab();

9.2.2.1. 初始化模型参数

接下来，我们需要定义和初始化模型参数。 如前所述，超参数numHiddens定义隐藏单元的数量。 我们按照标准差\(0.01\)的高斯分布初始化权重，并将偏置项设为\(0\)。

public static NDList getLSTMParams(int vocabSize, int numHiddens, Device device) {
    int numInputs = vocabSize;
    int numOutputs = vocabSize;

    // Input gate parameters
    NDList temp = three(numInputs, numHiddens, device);
    NDArray W_xi = temp.get(0);
    NDArray W_hi = temp.get(1);
    NDArray b_i = temp.get(2);

    // Forget gate parameters
    temp = three(numInputs, numHiddens, device);
    NDArray W_xf = temp.get(0);
    NDArray W_hf = temp.get(1);
    NDArray b_f = temp.get(2);

    // Output gate parameters
    temp = three(numInputs, numHiddens, device);
    NDArray W_xo = temp.get(0);
    NDArray W_ho = temp.get(1);
    NDArray b_o = temp.get(2);

    // Candidate memory cell parameters
    temp = three(numInputs, numHiddens, device);
    NDArray W_xc = temp.get(0);
    NDArray W_hc = temp.get(1);
    NDArray b_c = temp.get(2);

    // Output layer parameters
    NDArray W_hq = normal(new Shape(numHiddens, numOutputs), device);
    NDArray b_q = manager.zeros(new Shape(numOutputs), DataType.FLOAT32, device);

    // Attach gradients
    NDList params =
            new NDList(
                    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);
    for (NDArray param : params) {
        param.setRequiresGradient(true);
    }
    return params;
}

public static NDArray normal(Shape shape, Device device) {
    return manager.randomNormal(0, 0.01f, shape, DataType.FLOAT32, device);
}

public static NDList three(int numInputs, int numHiddens, Device device) {
    return new NDList(
            normal(new Shape(numInputs, numHiddens), device),
            normal(new Shape(numHiddens, numHiddens), device),
            manager.zeros(new Shape(numHiddens), DataType.FLOAT32, device));
}

9.2.2.2. 定义模型

在初始化函数中， 长短期记忆网络的隐状态需要返回一个额外的记忆元， 单元的值为0，形状为（批量大小，隐藏单元数）。 因此，我们得到以下的状态初始化。

public static NDList initLSTMState(int batchSize, int numHiddens, Device device) {
    return new NDList(
            manager.zeros(new Shape(batchSize, numHiddens), DataType.FLOAT32, device),
            manager.zeros(new Shape(batchSize, numHiddens), DataType.FLOAT32, device));
}

实际模型的定义与我们前面讨论的一样： 提供三个门和一个额外的记忆元。 请注意，只有隐状态才会传递到输出层， 而记忆元\(\mathbf{C}_t\)不直接参与输出计算。

public static Pair<NDArray, NDList> lstm(NDArray inputs, NDList state, NDList params) {
    NDArray W_xi = params.get(0);
    NDArray W_hi = params.get(1);
    NDArray b_i = params.get(2);

    NDArray W_xf = params.get(3);
    NDArray W_hf = params.get(4);
    NDArray b_f = params.get(5);

    NDArray W_xo = params.get(6);
    NDArray W_ho = params.get(7);
    NDArray b_o = params.get(8);

    NDArray W_xc = params.get(9);
    NDArray W_hc = params.get(10);
    NDArray b_c = params.get(11);

    NDArray W_hq = params.get(12);
    NDArray b_q = params.get(13);

    NDArray H = state.get(0);
    NDArray C = state.get(1);
    NDList outputs = new NDList();
    NDArray X, Y, I, F, O, C_tilda;
    for (int i = 0; i < inputs.size(0); i++) {
        X = inputs.get(i);
        I = Activation.sigmoid(X.dot(W_xi).add(H.dot(W_hi).add(b_i)));
        F = Activation.sigmoid(X.dot(W_xf).add(H.dot(W_hf).add(b_f)));
        O = Activation.sigmoid(X.dot(W_xo).add(H.dot(W_ho).add(b_o)));
        C_tilda = Activation.tanh(X.dot(W_xc).add(H.dot(W_hc).add(b_c)));
        C = F.mul(C).add(I.mul(C_tilda));
        H = O.mul(Activation.tanh(C));
        Y = H.dot(W_hq).add(b_q);
        outputs.add(Y);
    }
    return new Pair(
            outputs.size() > 1 ? NDArrays.concat(outputs) : outputs.get(0), new NDList(H, C));
}

9.2.2.3. 训练和预测

让我们通过实例化 Section 8.5中 引入的RNNModelScratch类来训练一个长短期记忆网络， 就如我们在 Section 9.1中所做的一样。

int vocabSize = vocab.length();
int numHiddens = 256;
Device device = manager.getDevice();
int numEpochs = Integer.getInteger("MAX_EPOCH", 500);

int lr = 1;

Functions.TriFunction<Integer, Integer, Device, NDList> getParamsFn =
        (a, b, c) -> getLSTMParams(a, b, c);
Functions.TriFunction<Integer, Integer, Device, NDList> initLSTMStateFn =
        (a, b, c) -> initLSTMState(a, b, c);
Functions.TriFunction<NDArray, NDList, NDList, Pair<NDArray, NDList>> lstmFn = (a, b, c) -> lstm(a, b, c);

RNNModelScratch model =
        new RNNModelScratch(
                vocabSize, numHiddens, device, getParamsFn, initLSTMStateFn, lstmFn);
TimeMachine.trainCh8(model, dataset, vocab, lr, numEpochs, device, false, manager);

perplexity: 1.1, 12066.0 tokens/sec on gpu(0)
time travellerit s against reason said filby but you willnever c
travellerype thing that by madeateryienclery is it huss ge

9.2.3. 简洁实现

使用高级API，我们可以直接实例化LSTM模型。 高级API封装了前文介绍的所有配置细节。 这段代码的运行速度要快得多， 因为它使用的是编译好的运算符而不是Python来处理之前阐述的许多细节。

LSTM lstmLayer =
        LSTM.builder()
                .setNumLayers(1)
                .setStateSize(numHiddens)
                .optReturnState(true)
                .optBatchFirst(false)
                .build();
RNNModel modelConcise = new RNNModel(lstmLayer, vocab.length());
TimeMachine.trainCh8(modelConcise, dataset, vocab, lr, numEpochs, device, false, manager);

INFO Training on: 1 GPUs.
INFO Load MXNet Engine Version 1.9.0 in 0.062 ms.

perplexity: 1.1, 79980.0 tokens/sec on gpu(0)
time traveller file some abeeimenthon this tore proner arspowsis
traveller fores yound at at sughtare mede a soit and said t

长短期记忆网络是典型的具有重要状态控制的隐变量自回归模型。 多年来已经提出了其许多变体，例如，多层、残差连接、不同类型的正则化。 然而，由于序列的长距离依赖性，训练长短期记忆网络 和其他序列模型（例如门控循环单元）的成本是相当高的。 在后面的内容中，我们将讲述更高级的替代模型，如transformer。

9.2.4. 小结

  • 长短期记忆网络有三种类型的门：输入门、遗忘门和输出门。

  • 长短期记忆网络的隐藏层输出包括“隐状态”和“记忆元”。只有隐状态会传递到输出层，而记忆元完全属于内部信息。

  • 长短期记忆网络可以缓解梯度消失和梯度爆炸。

9.2.5. 练习

  1. 调整和分析超参数对运行时间、困惑度和输出顺序的影响。

  2. 你需要如何更改模型以生成适当的单词，而不是字符序列？

  3. 在给定隐藏层维度的情况下，比较门控循环单元、长短期记忆网络和常规循环神经网络的计算成本。要特别注意训练和推断成本。

  4. 既然候选记忆元通过使用\(\tanh\)函数来确保值范围在\((-1,1)\)之间，那么为什么隐状态需要再次使用\(\tanh\)函数来确保输出值范围在\((-1,1)\)之间呢？

  5. 实现一个能够基于时间序列进行预测而不是基于字符序列进行预测的长短期记忆网络模型。