Understanding Recurrent Neural Networks Step-by-Step with PyTorch
Recurrent neural networks (RNNs) are a powerful class of neural networks that have revolutionized many fields involving sequential data, such as natural language processing, speech recognition, and time series forecasting. Their ability to capture long-term dependencies and maintain an internal memory has made them the go-to architecture for modeling sequences.
In this deep dive, we‘ll build a solid understanding of RNNs from the ground up and learn how to implement them in PyTorch. We‘ll start with the basic concepts and equations, build intuition through interactive demos, and walk through code step-by-step. By the end, you‘ll be well-equipped to apply RNNs to real-world sequence modeling problems.
A Brief History of RNNs
The story of RNNs begins in the 1980s with the seminal work of John Hopfield on recurrent neural networks that could learn and recall patterns [1]. These early Hopfield networks laid the foundation for future RNN architectures. However, training RNNs proved difficult due to the vanishing and exploding gradient problems first identified by Sepp Hochreiter in 1991 [2].
In 1997, Hochreiter and Jürgen Schmidhuber proposed Long Short-Term Memory (LSTM) networks to address these challenges [3]. LSTMs introduced gating mechanisms that allow the network to better control the flow of information and learn long-range dependencies. This breakthrough sparked a renaissance in RNNs and enabled many of the impressive applications we see today.
More recently, gated recurrent units (GRUs) were proposed by Cho et al. in 2014 as a simpler alternative to LSTMs [4]. In 2017, Vaswani et al. introduced the Transformer architecture which foregoes recurrence altogether in favor of attention mechanisms [5]. Transformers have since taken the NLP world by storm, but RNNs remain an important and widely-used tool.
Unrolling an RNN
At its core, an RNN is a neural network that processes a sequence of inputs one element at a time, maintaining an internal hidden state that encodes information about the past inputs. Let‘s first build some intuition by visually unrolling an RNN.
At each time step $t$, the RNN takes in an input $xt$ and the previous hidden state $h{t-1}$ and produces an output $o_t$ and a new hidden state $h_t$. The same weights $W$ are used at every time step, allowing the network to learn patterns across different positions in the sequence. This weight sharing is what allows RNNs to generalize to sequences of variable length.
Mathematically, the computations at each time step are as follows:
$$ht = \tanh(W{hx} xt + W{hh} h_{t-1} + b_h)$$
$$ot = W{oh} h_t + b_o$$
where $W{hx}$, $W{hh}$, and $W_{oh}$ are learnable weight matrices, $b_h$ and $b_o$ are bias vectors, and $\tanh$ is the hyperbolic tangent activation function.
The Vanishing and Exploding Gradient Problem
While the above equations might seem simple, training an RNN is quite challenging due to the vanishing and exploding gradient problem. During backpropagation, the gradients are multiplied by the recurrent weight matrix $W_{hh}$ at each time step. If the weights are small (less than 1), the gradients will exponentially shrink as they are propagated back in time, making it difficult to learn long-term dependencies. Conversely, if the weights are large, the gradients can exponentially grow and cause training to diverge.
This problem was a major obstacle in training RNNs until the advent of Long Short-Term Memory networks in 1997 [3]. LSTMs introduce gating mechanisms that allow the network to control the flow of information and better capture long-range dependencies.
An LSTM cell has three main gates:
- Forget gate: Controls what information to discard from the previous cell state
- Input gate: Controls what new information to add to the cell state
- Output gate: Controls what information from the cell state to output
These gates allow the LSTM to selectively remember or forget information over long sequences, mitigating the vanishing gradient problem. GRUs are a newer and simpler variant of LSTMs that combine the input and forget gates into a single update gate.
Implementing RNNs in PyTorch
Now that we have a conceptual understanding of RNNs, let‘s see how to implement them in PyTorch. We‘ll start with a basic Elman RNN and then show how to extend it to LSTMs and GRUs.
First, we define our RNN class:
class RNN(nn.Module):
def __init__(self, input_size, hidden_size, num_layers, num_classes):
super(RNN, self).__init__()
self.hidden_size = hidden_size
self.num_layers = num_layers
self.rnn = nn.RNN(input_size, hidden_size, num_layers, batch_first=True)
self.fc = nn.Linear(hidden_size, num_classes)
def forward(self, x):
h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)
out, _ = self.rnn(x, h0)
out = self.fc(out[:, -1, :])
return outThe key component is the nn.RNN module which takes in the input size, hidden size, and number of layers. We also define a fully-connected layer nn.Linear to map the final hidden state to the output classes.
In the forward pass, we first initialize the hidden state h0 to all zeros. We then pass the input x and initial hidden state to the RNN module which returns the output out at each time step and the final hidden state. We take the output at the last time step out[:, -1, :] and pass it through the fully-connected layer to get the final output.
To extend this to an LSTM, we simply replace nn.RNN with nn.LSTM:
self.rnn = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)and modify the forward pass to include the initial cell state c0:
def forward(self, x):
h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)
c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)
out, _ = self.rnn(x, (h0, c0))
out = self.fc(out[:, -1, :])
return outPyTorch‘s RNN modules expect the input data to be in (batch, seq_len, input_size) format if batch_first=True. Here‘s an example of preparing a toy sine wave dataset:
seq_length = 20
input_size = 1
num_samples = 1000
time_steps = np.linspace(0, np.pi, seq_length)
data = np.sin(time_steps)
data = data.reshape((num_samples, seq_length, input_size))We can then instantiate our RNN model and train it using standard PyTorch syntax:
model = RNN(input_size=1, hidden_size=16, num_layers=1, num_classes=1).to(device)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
for epoch in range(100):
for i in range(num_batches):
inputs, targets = get_batch(i, batch_size=32)
inputs, targets = inputs.to(device), targets.to(device)
outputs = model(inputs)
loss = criterion(outputs, targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()After training, we can visualize the hidden states of the RNN to gain insight into how it is processing the sequence. The hidden state at each time step can be thought of as a learned representation that encodes the relevant information from the past inputs.
In this example, we see that the hidden states evolve over time to capture the underlying sine wave pattern. This learned representation allows the RNN to make predictions based on the historical context.
Conclusion
Recurrent neural networks are a powerful tool for modeling sequential data across a wide range of domains. In this article, we demystified the inner workings of RNNs and showed how to implement them step-by-step in PyTorch.
We started with the basic concepts of unrolling an RNN and the key equations governing the hidden state and output at each time step. We then discussed the vanishing and exploding gradient problem and how gated architectures like LSTMs and GRUs address this challenge.
On the implementation side, we saw how to define an RNN in PyTorch using the nn.RNN module and how to prepare sequential data for training. We also visualized the learned hidden states to build intuition for how RNNs process and represent sequences.
Here are some key takeaways and best practices to keep in mind when working with RNNs:
- RNNs are well-suited for tasks involving sequential data where the order matters, such as time series, natural language, and speech.
- Vanilla RNNs suffer from vanishing and exploding gradients which limit their ability to capture long-term dependencies. LSTMs and GRUs introduce gating mechanisms to mitigate this problem.
- PyTorch provides high-level APIs like
nn.RNN,nn.LSTM, andnn.GRUthat make it easy to build and train recurrent models. - Monitoring the learned hidden states can give insight into how the RNN is representing and processing the sequence data.
- Regularization techniques like dropout and weight decay can help prevent overfitting, especially on smaller datasets.
- While RNNs have been surpassed by Transformers for many tasks, they remain a powerful and widely-used tool in the deep learning practitioner‘s toolkit.
I hope this deep dive has demystified RNNs and equipped you with the knowledge and code to apply them to your own projects. For further reading, I recommend the following resources:
- Understanding LSTM Networks by Chris Olah
- The Unreasonable Effectiveness of Recurrent Neural Networks by Andrej Karpathy
- Supervised Sequence Labelling with Recurrent Neural Networks by Alex Graves
As you continue your deep learning journey, keep an eye out for the latest advances in recurrent architectures and be sure to experiment with different models and hyperparameters to get the best results on your specific problem. Happy coding!
