934 KiB
import torch
from torch import nn
from d2l import torch as d2lGated Memory Cell
Each memory cell is equipped with an internal state
and a number of multiplicative gates that determine whether
(i) a given input should impact the internal state (the input gate),
(ii) the internal state should be flushed to 0 (the forget gate),
and (iii) the internal state of a given neuron
should be allowed to impact the cell's output (the output gate).
Gated Hidden State
The key distinction between vanilla RNNs and LSTMs is that the latter support gating of the hidden state. This means that we have dedicated mechanisms for when a hidden state should be updated and also for when it should be reset. These mechanisms are learned and they address the concerns listed above. For instance, if the first token is of great importance we will learn not to update the hidden state after the first observation. Likewise, we will learn to skip irrelevant temporary observations. Last, we will learn to reset the latent state whenever needed. We discuss this in detail below.
Input Gate, Forget Gate, and Output Gate
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 :numref:fig_lstm_0.
Three fully connected layers with sigmoid activation functions
compute the values of the input, forget, and output gates.
As a result of the sigmoid activation,
all values of the three gates
are in the range of (0, 1).
Additionally, we require an input node,
typically computed with a tanh activation function.
Intuitively, the input gate determines how much
of the input node's value should be added
to the current memory cell internal state.
The forget gate determines whether to keep
the current value of the memory or flush it.
And the output gate determines whether
the memory cell should influence the output
at the current time step.
🏷️fig_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:
\begin{aligned}
\mathbf{I}_t &= \sigma(\mathbf{X}_t \mathbf{W}_{\textrm{xi}} + \mathbf{H}_{t-1} \mathbf{W}_{\textrm{hi}} + \mathbf{b}_\textrm{i}),\\
\mathbf{F}_t &= \sigma(\mathbf{X}_t \mathbf{W}_{\textrm{xf}} + \mathbf{H}_{t-1} \mathbf{W}_{\textrm{hf}} + \mathbf{b}_\textrm{f}),\\
\mathbf{O}_t &= \sigma(\mathbf{X}_t \mathbf{W}_{\textrm{xo}} + \mathbf{H}_{t-1} \mathbf{W}_{\textrm{ho}} + \mathbf{b}_\textrm{o}),
\end{aligned}
where \mathbf{W}_{\textrm{xi}}, \mathbf{W}_{\textrm{xf}}, \mathbf{W}_{\textrm{xo}} \in \mathbb{R}^{d \times h} and \mathbf{W}_{\textrm{hi}}, \mathbf{W}_{\textrm{hf}}, \mathbf{W}_{\textrm{ho}} \in \mathbb{R}^{h \times h} are weight parameters
and \mathbf{b}_\textrm{i}, \mathbf{b}_\textrm{f}, \mathbf{b}_\textrm{o} \in \mathbb{R}^{1 \times h} are bias parameters.
Note that broadcasting
(see :numref:subsec_broadcasting)
is triggered during the summation.
We use sigmoid functions
(as introduced in :numref:sec_mlp)
to map the input values to the interval (0, 1).
Input Node
Next we design the memory cell.
Since we have not specified the action of the various gates yet,
we first introduce the input node
\tilde{\mathbf{C}}_t \in \mathbb{R}^{n \times h}.
Its computation is similar to that of the three gates described above,
but uses a \tanh function with a value range for (-1, 1) as the activation function.
This leads to the following equation at time step t:
\tilde{\mathbf{C}}_t = \textrm{tanh}(\mathbf{X}_t \mathbf{W}_{\textrm{xc}} + \mathbf{H}_{t-1} \mathbf{W}_{\textrm{hc}} + \mathbf{b}_\textrm{c}),
where \mathbf{W}_{\textrm{xc}} \in \mathbb{R}^{d \times h} and \mathbf{W}_{\textrm{hc}} \in \mathbb{R}^{h \times h} are weight parameters and \mathbf{b}_\textrm{c} \in \mathbb{R}^{1 \times h} is a bias parameter.
A quick illustration of the input node is shown in :numref:fig_lstm_1.
🏷️fig_lstm_1
Memory Cell Internal State
In LSTMs, 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 cell internal state \mathbf{C}_{t-1} \in \mathbb{R}^{n \times h} we retain.
Using the Hadamard (elementwise) product operator \odot
we arrive at the following update equation:
\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 1 and the input gate is always 0,
the memory cell internal state \mathbf{C}_{t-1}
will remain constant forever,
passing unchanged to each subsequent time step.
However, input gates and forget gates
give the model the flexibility of being able to learn
when to keep this value unchanged
and when to perturb it in response
to subsequent inputs.
In practice, this design alleviates the vanishing gradient problem,
resulting in models that are much easier to train,
especially when facing datasets with long sequence lengths.
We thus arrive at the flow diagram in :numref:fig_lstm_2.
🏷️fig_lstm_2
Hidden State
Last, we need to define how to compute the output
of the memory cell, i.e., the hidden state \mathbf{H}_t \in \mathbb{R}^{n \times h}, as seen by other layers.
This is where the output gate comes into play.
In LSTMs, we first apply \tanh to the memory cell internal state
and then apply another point-wise multiplication,
this time with the output gate.
This ensures that the values of \mathbf{H}_t
are always in the interval (-1, 1):
\mathbf{H}_t = \mathbf{O}_t \odot \tanh(\mathbf{C}_t).
Whenever the output gate is close to 1,
we allow the memory cell internal state to impact the subsequent layers uninhibited,
whereas for output gate values close to 0,
we prevent the current memory from impacting other layers of the network
at the current time step.
Note that a memory cell can accrue information
across many time steps without impacting the rest of the network
(as long as the output gate takes values close to 0),
and then suddenly impact the network at a subsequent time step
as soon as the output gate flips from values close to 0
to values close to 1. :numref:fig_lstm_3 has a graphical illustration of the data flow.
🏷️fig_lstm_3
Implementation from Scratch
Now let's implement an LSTM from scratch.
As same as the experiments in :numref:sec_rnn-scratch,
we first load The Time Machine dataset.
[Initializing Model Parameters]
Next, we need to define and initialize the model parameters.
As previously, the hyperparameter num_hiddens
dictates the number of hidden units.
We initialize weights following a Gaussian distribution
with 0.01 standard deviation,
and we set the biases to 0.
class LSTMScratch(d2l.Module):
def __init__(self, num_inputs, num_hiddens, sigma=0.01):
super().__init__()
self.save_hyperparameters()
init_weight = lambda *shape: nn.Parameter(torch.randn(*shape) * sigma)
triple = lambda: (init_weight(num_inputs, num_hiddens),
init_weight(num_hiddens, num_hiddens),
nn.Parameter(torch.zeros(num_hiddens)))
self.W_xi, self.W_hi, self.b_i = triple() # Input gate
self.W_xf, self.W_hf, self.b_f = triple() # Forget gate
self.W_xo, self.W_ho, self.b_o = triple() # Output gate
self.W_xc, self.W_hc, self.b_c = triple() # Input node@d2l.add_to_class(LSTMScratch)
def forward(self, inputs, H_C=None):
if H_C is None:
# Initial state with shape: (batch_size, num_hiddens)
H = torch.zeros((inputs.shape[1], self.num_hiddens),
device=inputs.device)
C = torch.zeros((inputs.shape[1], self.num_hiddens),
device=inputs.device)
else:
H, C = H_C
outputs = []
for X in inputs:
I = torch.sigmoid(torch.matmul(X, self.W_xi) +
torch.matmul(H, self.W_hi) + self.b_i)
F = torch.sigmoid(torch.matmul(X, self.W_xf) +
torch.matmul(H, self.W_hf) + self.b_f)
O = torch.sigmoid(torch.matmul(X, self.W_xo) +
torch.matmul(H, self.W_ho) + self.b_o)
C_tilde = torch.tanh(torch.matmul(X, self.W_xc) +
torch.matmul(H, self.W_hc) + self.b_c)
C = F * C + I * C_tilde
H = O * torch.tanh(C)
outputs.append(H)
return outputs, (H, C)data = d2l.TimeMachine(batch_size=1024, num_steps=32)
lstm = LSTMScratch(num_inputs=len(data.vocab), num_hiddens=32)
model = d2l.RNNLMScratch(lstm, vocab_size=len(data.vocab), lr=4)
trainer = d2l.Trainer(max_epochs=50, gradient_clip_val=1, num_gpus=1)
trainer.fit(model, data)
class LSTM(d2l.RNN):
def __init__(self, num_inputs, num_hiddens):
d2l.Module.__init__(self)
self.save_hyperparameters()
self.rnn = nn.LSTM(num_inputs, num_hiddens)
def forward(self, inputs, H_C=None):
return self.rnn(inputs, H_C)lstm = LSTM(num_inputs=len(data.vocab), num_hiddens=32)
model = d2l.RNNLM(lstm, vocab_size=len(data.vocab), lr=4)
trainer.fit(model, data)
model.predict('it has', 20, data.vocab, d2l.try_gpu())'it has a the time travelly'