In the previous exercise you implemented a vanilla RNN and applied it to image captioning. In this notebook you will implement the LSTM update rule and use it for image captioning.

LSTM

If you read recent papers, you'll see that many people use a variant on the vanilla RNN called Long-Short Term Memory (LSTM) RNNs. Vanilla RNNs can be tough to train on long sequences due to vanishing and exploding gradients caused by repeated matrix multiplication. LSTMs solve this problem by replacing the simple update rule of the vanilla RNN with a gating mechanism as follows.

Similar to the vanilla RNN, at each timestep we receive an input π‘₯π‘‘βˆˆβ„π· and the previous hidden state β„Žπ‘‘βˆ’1βˆˆβ„π»; the LSTM also maintains an 𝐻-dimensional cell state, so we also receive the previous cell state π‘π‘‘βˆ’1βˆˆβ„π». The learnable parameters of the LSTM are an input-to-hidden matrix π‘Šπ‘₯βˆˆβ„4𝐻×𝐷, a hidden-to-hidden matrix π‘Šβ„Žβˆˆβ„4𝐻×𝐻 and a bias vector π‘βˆˆβ„4𝐻.

At each timestep we first compute an activation vector π‘Žβˆˆβ„4𝐻 as π‘Ž=π‘Šπ‘₯π‘₯𝑑+π‘Šβ„Žβ„Žπ‘‘βˆ’1+𝑏. We then divide this into four vectors π‘Žπ‘–,π‘Žπ‘“,π‘Žπ‘œ,π‘Žπ‘”βˆˆβ„π» where π‘Žπ‘– consists of the first 𝐻 elements of π‘Ž, π‘Žπ‘“ is the next 𝐻 elements of π‘Ž, etc. We then compute the input gate π‘”βˆˆβ„π», forget gate π‘“βˆˆβ„π», output gate π‘œβˆˆβ„π» and block input π‘”βˆˆβ„π» as

𝑖=𝜎(π‘Žπ‘–)𝑓=𝜎(π‘Žπ‘“)π‘œ=𝜎(π‘Žπ‘œ)𝑔=tanh(π‘Žπ‘”)

where 𝜎 is the sigmoid function and tanh is the hyperbolic tangent, both applied elementwise.

Finally we compute the next cell state 𝑐𝑑 and next hidden state β„Žπ‘‘ as

𝑐𝑑=π‘“βŠ™π‘π‘‘βˆ’1+π‘–βŠ™π‘”β„Žπ‘‘=π‘œβŠ™tanh(𝑐𝑑)

where βŠ™ is the elementwise product of vectors.

In the rest of the notebook we will implement the LSTM update rule and apply it to the image captioning task.