Implementation Differences in LSTM Layers: TensorFlow vs PyTorch | by Madhushan Buwaneswaran | Towards Data Science

Tensorflow and Pytorch are the two most widely used libraries in deep learning. Both these libraries have different approaches when it comes to implementing neural networks. Both these libraries started with major differences. Usually, TF was more preferred for developing production-ready models due to its better optimizations, and Pytorch was preferred for research works due to its more “pythonic” syntax and eager execution. But with Torchscript and TF 2.0, the gap between the two libraries has shrunk.
Both of these libraries have good community support and active contributions. So we can easily implement different types of neural networks without any major issues in both these libraries. But these two libraries have architectural differences. Hence, there are certain differences we need to focus on when implementing a neural network.
One such difference is in the layers API. A neural network (NN) can be designed using either the sub-classing approach or sequential model API approach in both these libraries. The sub-classing approach is more preferred out of the two since it is object-oriented and extensible. When implementing a NN in either of these libraries, we can use already designed layers — linear (fully connected) layers, convolution layers, recurrent layers, etc. and extend them into our model. In TF, we use tensorflow.keras.layers, and in Pytorch, we use torch.nn to access these layers. As I mentioned earlier, there are subtle differences in the implementation of these layers. Most of the time, they are minor and intuitive. But in LSTM (Long Short-Term Memory) layers, these differences are somewhat major and significant. The way the layer is parameterized, the default values for the parameters, and the default output of the layer mostly differ between Pytorch and TF in LSTM layers.
In this post, I will try to explain the differences between LSTM layers in these two libraries. I hope it will be of help if someone refers to a NN (with LSTM layers) implemented in one library and tries to reproduce it in the other.
I have added this section only for the sake of completion and as a refresher (in case one needs it). But if you are trying to understand the implementation differences when using LSTM layers, then I hope you already have a background in deep learning and know the fundamentals of LSTMs. So you can skip this part.
A recurrent neural network (RNN) is a special type of NN which is designed to learn from sequential data. A conventional NN would take an input and give a prediction based on that input only. It doesn’t look into previous inputs/outputs and make a decision. But when predicting the next value in sequential data (eg. — sentences, the daily average temperature in a city), we cannot do it by looking at the current data point only. We have to look at the behavior of the series (up to the current point) as a whole and get the ‘context of the series’ to make a meaningful prediction. This is what RNNs are specialized for. RNNs are used for tasks like language modeling, machine translation, sequence prediction, etc.
Fig. 1 shows how an RNN takes the current value and previous lateral output as input. The lateral output of the previous timestep is passed by the recurrent mechanism (as at). It is expected (and proven) that the RNN layer will learn to capture the ‘context’ of the sequence in the lateral output. Please note that the unrolled version of the RNN shown in fig. 1 (b) depicts the same layer as copies for better understanding and doesn’t represent multiple layers.
xt - Primary input at timestep t. Shape - (nx,1)
at - Lateral output at timestep t (also the lateral input with a unit time delay). Shape - (na,1)
ht - Primary output at timestep t. Shape - (nh,1)
Waa - Weight matrix that is multiplied with lateral input. Shape - (na,na)
Wax - Weight matrix that is multiplied with primary input. Shape - (na,nx)
ba - Bias term added when generating at. Shape - (na,1)
Wha - Weight matrix that is multiplied with at when generating ht. Shape - (nh,na)
bh - Bias term added when generating ht. Shape - (nh,1)
g() -  non linear activation function (eg.-  σ(), tanh() )
The design principle of RNN expects that the ‘context’ derived from earlier values in the series persists along the sequence for a long range. But it was observed in practical scenarios, the RNN’s ability to learn the ‘context’ of a sequence deteriorates with distance. That is, RNNs are not good enough to capture long-term dependencies in sequences. Long Short-Term Memory (LSTM) networks are a specialized version of RNN that was introduced with the intention of preserving long-term ‘context’ information in sequences. The inherent design of LSTMs enables them to capture the long-term context through what is known as cell state (denoted by ct). The diagram below (Fig. 2) shows a typical LSTM layer. The diagram I used is based on the diagrams used in a blog post on understanding LSTMs by Chris Olah. (It is a wonderful blog post. Do check it out.)
xt - Primary input at timestep t. Shape - (nx,1)
ct - Cell state at timestep t. Shape - (nh,1)
ht - Primary output at timestep t. Shape - (nh,1)
c~t - Candidate vector to update ct. Shape - (nh,1)
Γf - Forget gate mask. Shape - (nh,1)
Γu - Update gate mask. Shape - (nh,1)
Γo - Output gate mask. Shape - (nh,1)
Wc - Weight matrix used to generate candidate vector. Shape - (nh,nh+nx)
bc - Bias term used to generate candidate vector. Shape (nh,1)
Wf - Weight matrix in forget gate. Shape - (nh,nh+nx)
bf - Bias term in forget gate. Shape - (nh,1)
Wu - Weight matrix in update gate. Shape - (nh,nh+nx)
bu - Bias term in update gate. Shape - (nh,1)
Wo - Weight matrix in output gate. Shape - (nh,nh+nx)
bo - Bias term in output gate. Shape - (nh,1)
As Chris explains in his blog, the cell state can be thought of as a conveyor belt, which carries the long-term ‘context’ from start to end of the sequence. At each timestep based on the current input (xt) and the previous output (ht−1), the LSTM layer decides what information to forget from the ‘context’ and what information to add to the ‘context’. The forgetting is handled by the forget gate (Γf). Γf can be thought of as a mask that is applied to ct−1. If a value in a certain position in Γf is 0 (or closer), when multiplied with ct−1, it removes the information in that position from the context. If the value of the mask is 1 (or closer) in some other position, after multiplication with ct−1, it allows the information in that position to persist in the context. The operation Γf.*ct−1 (.* — element-wise multiplication) makes sure that the unwanted part of the context is forgotten. The addition of new information to the cell state is handled in two steps. First, a candidate vector (c̃t) is created through the sub-layer with tanh activation, based on xt and ht−1. Then the update gate (Γu) generates a mask (just like in forget gate) which decides what part of the candidate vector is added to the cell state. ThenΓu.*c̃t is added to ct−1 to generate ct. Finally, to generate the output, we pass ct through a tanh non-linearity and use it. But to decide what part of it to give out, we make use of the update gate (Γo). Γo.*tanh(ct) is given out as the output of the LSTM layer at timestep t.
Please note that both ct−1 and ht−1 are given as lateral inputs in an LSTM layer, compared to the vanilla RNN layer where only at−1 is given.
If you need additional information on the principles of RNN and LSTM, I would recommend going through the first week of the Sequence Models course by Andrew Ng (it is free to audit) and reading this wonderful blog post by Chris Olah.
At the time of writing Tensorflow version was 2.4.1
In TF, we can use tf.keras.layers.LSTM and create an LSTM layer. When initializing an LSTM layer, the only required parameter is units. The parameter units corresponds to the number of output features of that layer. That is units = nh in our terminology. nx will be inferred from the output of the previous layer. Hence the library can initialize all the weight and bias terms in the LSTM layer.
TF LSTM layer expects a 3 dimensional tensor as input during forward propagation. This input should be of the shape (batch, timesteps, input_features). This is shown in the code snippet below. Suppose we are using this LSTM layer to train a language model. Our input will be sentences. The first dimension corresponds to how many sentences we use as one batch to train the model. The second dimension corresponds to how many words are present in one such sentence. In practical setting, the number of words in each sentence varies from sentence to sentence. So, in order to batch these sentences, we can select the length of the longest sentence in the training corpus as this dimension and pad the other sentences with trailing zeros. The last dimension corresponds to the number of features used to represent each word. For simplicity, if we say, we are using one-hot encoding and there are 10000 words in our vocabulary, then this dimension will be 10000.
import tensorflow as tf

batch_size = 32
timesteps = 10
input_features = 8
output_features = 4

inputs = tf.random.normal([batch_size, timesteps, input_features])
lstm = tf.keras.layers.LSTM(units=output_features)
output = lstm(inputs)

print("Output shape of the LSTM layer's default output : ", output.shape)

Output shape of the LSTM layer's default output :  (32, 4)

The output is of the shape :    (batch_size, output_features) - (nbatch,nh)
This output corresponds to hT where T is the last timestep in the sequence.
But when initializing the layer, if we set time_major = True, then the input will be expected in the shape - (timesteps, batch, feature).
As seen from the above code snippet, the output of the LSTM (with default parameters) is of shape (32,4), which corresponds to (batch, output_features). So if we go back to the example of the language model, the output has one vector per sentence, with nh number of features per sentence (nh = units = no. of output features). This one vector (per sentence) is the output of the LSTM layer corresponding to the last timestep T (last word of the sentence). This output is hT in our notation. This is depicted in fig. 3.
However, if we want to stack multiple LSTM layers, the next layer also expects a time series as input. In such situations, we can set return_sequences=True when initiating the layer. Then the output will be of shape (32,10,4), which corresponds to (batch, timesteps, output_features). If return_sequence is set to True, then ht : ∀t = 1,2...T will be returned as output. This is shown in the code snippet below and the first LSTM layer in fig. 4.
If we want to get the cell state (ct) as an output, we need to set return_state=True when initiating the layer. Then We get a list of 3 tensors as output. According to documentation, if we set both return_sequences=True and return_state=True, these three tensors will be — whole_seq_output, final_memory_state,and final_carry_state. This is shown in the code snippet below.
In our notation,
whole_seq_output — output corresponding to all timesteps. ht : ∀t = 1,2...T ; Shape — (batch, timesteps, output_features)
final_memory_state — Output corresponding to the last timestep.hT; Shape — (batch, output_features)
final_carry_state — Last cell state.cT; Shape — (batch, output_features)
If we set return_sequences=False and return_state=True, then these three tensors will be — final_memory_state, final_memory_state, and final_carry_state.
A single LSTM layer has five places where activation functions are used. But if we look at the parameters, we see only two parameters to set activation functions — activation and recurrent_activation. If we set a value to the activation parameter, it changes the activation applied to the candidate vector and the activation applied to the cell state just before element-wise multiplication with the output gate. Setting a value to recurrent_activation will change the activation functions of forget gate, update gate, and output gate.
Other parameters are quite self-explanatory or seldom used. One other thing to note is, we can set unroll=True, and the network will be unrolled. It will speed up the training process but will be memory intensive (since the same layer is copied multiple times).
The following code snippet implements the model shown in fig. 4 using TF. Note the output shape of each layer and the number of trainable parameters in each layer.
At the time of writing, Pytorch version was 1.8.1
In Pytorch, an LSTM layer can be created using torch.nn.LSTM. It requires two parameters at initiation input_size and hidden_size.input_size and hidden_size correspond to the number of input features to the layer and the number of output features of that layer, respectively. In our terminology, hidden_size = nh and input_size = nx.
During forward propagation, the Pytorch LSTM layer expects a 3 dimensional tensor as input (similar to TF). But the default ordering of the dimensions changes. The input tensor should be of shape (timesteps, batch, input_features). If we want to get the same order of dimensions as TF, we should set batch_first=True at layer initiation.
Another major difference that can be seen in Pytorch LSTM API is that, at initiation, we can set num_layers=k and initiate a block of k LSTM layers stacked as a single object. However, I personally do not prefer this approach since it makes the overall implementation less readable and less maintainable.
The next big difference is the output of the Pytorch LSTM layer. The output of the Pytorch LSTM layer is a tuple with two elements. The first element of the tuple is LSTM’s output corresponding to all timesteps (ht : ∀t = 1,2...T) with shape (timesteps, batch, output_features). The second element of the tuple is another tuple with two elements. The first element of this second tuple is the output corresponding to the last timestep (hT). It has the shape (1, batch, output_features). The second element of this second tuple is the cell state corresponding to the last timestep (cT). It also has the shape (1, batch, output_features). If we had initiated the LSTM as a block of stacked layers by setting num_layers=k, then hT and cT would have the shape (k, batch, output_features). Here, both hT and cT would have the last states of all the k layers in the stack. Further at initiation, had we set batch_first=True, then the timesteps and batch dimensions would swap in the output (similar to the input).
As far as I know, changing the activation functions inside the LSTM layer is not possible in Pytorch. Also, it is not possible to limit the LSTM layer to give only one output (as in TF). However, we can assign the outputs to variables, use the required output, and ignore the others, as shown in the code segment below. Apart from that, other parameters should be self-explanatory.
Note that if you set the LSTM layer to be bidirectional (which I didn’t talk about in this article), then the output shape will be different from what I have mentioned above. Please refer to the documentation for such a scenario.
The following code snippet implements the model shown in fig. 4 using Pytorch. Note the output shape of each layer and the number of trainable parameters in each layer.
If we look at the number of parameters in the two implementations of the model in fig. 4., it can be observed that there is a difference in the number of parameters in LSTM layers. This is again a design choice. Pytorch does a minor change when implementing the LSTM equations (1), (2), (3), and (4). TF adds a single bias vector (as in our equations) in each of these equations. But Pytorch (as shown here) adds two bias vectors per equation. Since the shape of each bias vector is (nh,1) and four such additional vectors are there per layer, there will be 4*nh more parameters in the Pytorch LSTM layers. In LSTM1 layer nh=8, so there are 32 additional parameters. In LSTM2 layer nh=4, so there are 16 more parameters.