Run this notebook online: or Colab:

9.2. Long Short-Term Memory (LSTM)

The challenge to address long-term information preservation and short-term input skipping in latent variable models has existed for a long time. One of the earliest approaches to address this was the long short-term memory (LSTM) [Hochreiter & Schmidhuber, 1997]. It shares many of the properties of the GRU. Interestingly, LSTMs have a slightly more complex design than GRUs but predates GRUs by almost two decades.

9.2.1. Gated Memory Cell

Arguably LSTM’s design is inspired by logic gates of a computer. LSTM introduces a memory cell (or cell for short) that has the same shape as the hidden state (some literatures consider the memory cell as a special type of the hidden state), engineered to record additional information. To control the memory cell we need a number of gates. One gate is needed to read out the entries from the cell. We will refer to this as the output gate. A second gate is needed to decide when to read data into the cell. We refer to this as the input gate. Last, we need a mechanism to reset the content of the cell, governed by a forget gate. The motivation for such a design is the same as that of GRUs, namely to be able to decide when to remember and when to ignore inputs in the hidden state via a dedicated mechanism. Let us see how this works in practice.

9.2.1.1. Input Gate, Forget Gate, and Output Gate

Just like in GRUs, the data feeding into the LSTM gates are the input at the current time step and the hidden state of the previous time step, as illustrated in lstm_0. They are processed by three fully-connected layers with a sigmoid activation function to compute the values of the input, forget. and output gates. As a result, values of the three gates are in the range of \((0, 1)\).

.. _lstm_0:

Mathematically, suppose that there are \(h\) hidden units, the batch size is \(n\), and the number of inputs is \(d\). Thus, the input is \(\mathbf{X}_t \in \mathbb{R}^{n \times d}\) and the hidden state of the previous time step is \(\mathbf{H}_{t-1} \in \mathbb{R}^{n \times h}\). Correspondingly, the gates at time step \(t\) are defined as follows: the input gate is \(\mathbf{I}_t \in \mathbb{R}^{n \times h}\), the forget gate is \(\mathbf{F}_t \in \mathbb{R}^{n \times h}\), and the output gate is \(\mathbf{O}_t \in \mathbb{R}^{n \times h}\). They are calculated as follows:

(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}\]

where \(\mathbf{W}_{xi}, \mathbf{W}_{xf}, \mathbf{W}_{xo} \in \mathbb{R}^{d \times h}\) and \(\mathbf{W}_{hi}, \mathbf{W}_{hf}, \mathbf{W}_{ho} \in \mathbb{R}^{h \times h}\) are weight parameters and \(\mathbf{b}_i, \mathbf{b}_f, \mathbf{b}_o \in \mathbb{R}^{1 \times h}\) are bias parameters.

9.2.1.2. Candidate Memory Cell

Next we design the memory cell. Since we have not specified the action of the various gates yet, we first introduce the candidate memory cell \(\tilde{\mathbf{C}}_t \in \mathbb{R}^{n \times h}\). Its computation is similar to that of the three gates described above, but using a \(\tanh\) function with a value range for \((-1, 1)\) as the activation function. This leads to the following equation at time step \(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),\]

where \(\mathbf{W}_{xc} \in \mathbb{R}^{d \times h}\) and \(\mathbf{W}_{hc} \in \mathbb{R}^{h \times h}\) are weight parameters and \(\mathbf{b}_c \in \mathbb{R}^{1 \times h}\) is a bias parameter.

A quick illustration of the candidate memory cell is shown in Section 9.2.1.2.

9.2.1.3. Memory Cell

In GRUs, we have a mechanism to govern input and forgetting (or skipping). Similarly, in LSTMs we have two dedicated gates for such purposes: the input gate \(\mathbf{I}_t\) governs how much we take new data into account via \(\tilde{\mathbf{C}}_t\) and the forget gate \(\mathbf{F}_t\) addresses how much of the old memory cell content \(\mathbf{C}_{t-1} \in \mathbb{R}^{n \times h}\) we retain. Using the same pointwise multiplication trick as before, we arrive at the following update equation:

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

If the forget gate is always approximately 1 and the input gate is always approximately 0, the past memory cells \(\mathbf{C}_{t-1}\) will be saved over time and passed to the current time step. This design is introduced to alleviate the vanishing gradient problem and to better capture long range dependencies within sequences.

We thus arrive at the flow diagram in Fig. 9.2.1.

Fig. 9.2.1 Computing the memory cell in an LSTM model.

9.2.1.4. Hidden State

Last, we need to define how to compute the hidden state \(\mathbf{H}_t \in \mathbb{R}^{n \times h}\). This is where the output gate comes into play. In LSTM it is simply a gated version of the \(\tanh\) of the memory cell. This ensures that the values of \(\mathbf{H}_t\) are always in the interval \((-1, 1)\).

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

Whenever the output gate approximates 1 we effectively pass all memory information through to the predictor, whereas for the output gate close to 0 we retain all the information only within the memory cell and perform no further processing.

lstm_3 has a graphical illustration of the data flow.

.. _lstm_3:

9.2.2. Implementation from Scratch

Now let us implement an LSTM from scratch. As same as the experiments in Section 8.5, we first load the time machine dataset.

%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. Initializing Model Parameters

Next we need to define and initialize the model parameters. As previously, the hyperparameter numHiddens defines the number of hidden units. We initialize weights following a Gaussian distribution with 0.01 standard deviation, and we set the biases to 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. Defining the Model

In the initialization function, the hidden state of the LSTM needs to return an additional memory cell with a value of 0 and a shape of (batch size, number of hidden units). Hence we get the following state initialization.

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));
}

The actual model is defined just like what we discussed before: providing three gates and an auxiliary memory cell. Note that only the hidden state is passed to the output layer. The memory cell \(\mathbf{C}_t\) does not directly participate in the output computation.

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. Training and Prediction

Let us train an LSTM as same as what we did in Section 9.1, by instantiating the RNNModelScratch class as introduced in Section 8.5.

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, 11571.5 tokens/sec on gpu(0)
time traveller a deald thas is all the earthat that ir mist a ve
traveller after the pauserequired for the proper assimilati

9.2.3. Concise Implementation

Using high-level APIs, we can directly instantiate an LSTM model. This encapsulates all the configuration details that we made explicit above. The code is significantly faster as it uses compiled operators rather than Java for many details that we spelled out in detail before.

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.055 ms.

perplexity: 1.1, 80819.7 tokens/sec on gpu(0)
time traveller for so it will be convenient to speak of himwas e
traveller frece thef and some there there wermat of lyon ab

LSTMs are the prototypical latent variable autoregressive model with nontrivial state control. Many variants thereof have been proposed over the years, e.g., multiple layers, residual connections, different types of regularization. However, training LSTMs and other sequence models (such as GRUs) are quite costly due to the long range dependency of the sequence. Later we will encounter alternative models such as Transformers that can be used in some cases.

9.2.4. Summary

• LSTMs have three types of gates: input gates, forget gates, and output gates that control the flow of information.

• The hidden layer output of LSTM includes the hidden state and the memory cell. Only the hidden state is passed into the output layer. The memory cell is entirely internal.

• LSTMs can alleviate vanishing and exploding gradients.

9.2.5. Exercises

1. Adjust the hyperparameters and analyze the their influence on running time, perplexity, and the output sequence.

2. How would you need to change the model to generate proper words as opposed to sequences of characters?

3. Compare the computational cost for GRUs, LSTMs, and regular RNNs for a given hidden dimension. Pay special attention to the training and inference cost.

4. Since the candidate memory cell ensures that the value range is between \(-1\) and \(1\) by using the \(\tanh\) function, why does the hidden state need to use the \(\tanh\) function again to ensure that the output value range is between \(-1\) and \(1\)?

5. Implement an LSTM model for time series prediction rather than character sequence prediction.