Long Short – Term Memory (LSTM)
Recommended read:
Standard RNNs are great at handling sequence data, but they suffer from a notoriously short attention span. As new information flows in, the older information is continuously overwritten and lost.
Imagine feeding an RNN this story: “The knight is brave. He fights using his ___ “
For our RNN model to correctly predict “sword”, the model needs the context from the word “knight”. As you already know, RNN uses BPTT to “learn”. For this scenario, let’s say that the local gradient is 0.5. Let’s trace the backpropagation (BPTT).
| Time Step (t) | Word | Local Gradient | Backward Math | Error Signal Size | Impact on Learning |
|---|---|---|---|---|---|
| 9 | ___ (Target) | — | Starting Error | 1.0 (100%) | Error is captured |
| 8 | his | 0.5 | 1.0 \times 0.5 | 0.5 (50%) | Strong |
| 7 | using | 0.5 | 0.5 \times 0.5 | 0.25 (25%) | Moderate |
| 6 | fights | 0.5 | 0.25 \times 0.5 | 0.125 (12.5%) | Fading |
| 5 | He | 0.5 | 0.125 \times 0.5 | 0.0625 (6.25%) | Weak |
| 4 | brave. | 0.5 | 0.0625 \times 0.5 | 0.03125 (3.1%) | Very Weak |
| 3 | is | 0.5 | 0.03125 \times 0.5 | 0.015625 (1.5%) | Barely registers |
| 2 | knight | 0.5 | 0.015625 \times 0.5 | 0.0078125 (0.78%) | Virtually None! |
| 1 | The | 0.5 | 0.0078125 \times 0.5 | 0.00390625 (0.39%) | Ignored |
By the time the error signal crawled back to the word “knight”, which is an important context for it to guess the weapon, it’s already been diminished to less than 1%. Since the signal is virtually invisible, the network cannot update its weights to connect ‘using his __ ‘ (predicting weapon) with the critical context word “knight”. This problem is called the Vanishing Gradient problem.
Intuition (Standard RNN vs LSTM)
Before we go to the actual math, it’ll be beneficial for us if we have a good intuition of what LSTM is compared to a standard RNN.
Standard RNN (Single-Slot Inventory)
Imagine playing an entire game where your character only has one single inventory slot. Every time you find a new item, you can trash the current one and replace it with the new item, or you can melt and fuse the old and the new item. By the time you reach the final boss, your early-game item are completely lost (most likely you replaced it with stronger, end-game item). The system is made so that the latest item you pick up is the most relevant.
LSTM (Backpack and Action Hotbar)
Now imagine instead of single slot, your character has a Backpack (Long-term memory), and a quick-access Action Hotbar (Short-term memory).
When you buy an item from an NPC or defeat an enemy and it drops an item, you actively manage your items and gears:
- Drop junk (Forget Gate) – you look at your character’s current environment, open your backpack and check your Action hotbar, then decides which items are not needed in the current circumstance (low-level items, finished quest items) then dispose them.
- Move items from Action Hotbar to backpack + Loot new items (Input Gate) – you look at your character’s current environment, you look at your Action Hotbar and the new loot items then decides which items you want to store in the backpack.
- Equip / Update action hotbar (Output Gate) – not all items in your backpack are needed immediately, but there are some items that you need to be easily accessible/usable (like HP/SP pots or elemental converters) based on your character’s current environment.
Because the items are safely stored in your backpack, you can still have an item you picked up from when you are Level 1 all the way to Level 99. Meanwhile, you are constantly changing the items in your Action Hotbar depending on the situation, for example if you are in a fire dungeon, you have water elemental converter in your action hotbar. Then when you move to a grass dungeon, you replace the water elemental converter with a fire elemental converter.
If we create a diagram on how our character drops junk items, loot new items, store items in inventory, and updates the action hotbar. It will look something like this.
LSTM Unit
Hopefully the example and diagram above help us have better intuition of what happens inside an LSTM unit. In this section, we will now look into the actual mathematical equations and operations.
The data going in the LSTM unit are the input at the current time step $X_t$ and the hidden state of the previous time step $H_{t-1}$. As you can see in the diagram above, all three gates uses $\sigma$ sigmoid activation function to compute their values. As the result of the sigma activation, all values of the three gates are in the range of (0, 1).
Let’s look at each of the gates closer:
Forget gate $F_t$
Imagine your character needs to train in a high-level volcano dungeon (fire environment). You open your backpack and go over each of the items and see if they are relevant. For example this is the contents of your backpack.
| Item Name | Should Keep? |
|---|---|
| Low-level trash loot | No (0.0) |
| Water-Element Greatsword | Yes (1.0) |
| Low-level legendary accessory | Yes (1.0) |
| Low-level legendary weapon | Mostly Yes (0.7) |
In your backpack, you see some low-level loots that you picked up in the past when you are still low-level, and you decide that it’s no longer relevant and just taking space, so you want to drop/discard it.
Now you see a Water-Element Greatsword. You know that this item will be your primary weapon in this volcano dungeon, so of course you want to keep it.
Next, you look at the Low-level legendary accessory and weapon. While both of these items are low-level, they are legendary so you may want to keep them rather than disposing.
The forget gate decides how much information from the previous time step to retain or discard. A value close to 0 means that the LSTM decides that the information from the previous time step is no longer relevant and should be forgotten. While a value close to 1 means that the information from the previous time step is very much relevant and we need to keep it.
Input gate $I_t$ and the Candidate
Now imagine your character defeats the dungeon boss, and boss drops the following loots:
| Boss loot name | Should pickup? |
|---|---|
| High-level legendary weapon | Yes (1.0) |
| Some random boss drops | Yes (0.8) |
| Healing potion | No (0.0) |
Of course, you want to pick up the high-level legendary weapon drop, it’s why you entered this dungeon in the first place! As for the random boss drops, you also want to pick them up as they can be used to craft epic items. As for the Healing potion, you have no interest in picking it up (you have lots!).
Since the boss fight is also done, not only will you decide which items to loot, you will also check your current action hotbar and see which its items can be stored to the backpack.
The input gate determines how much of the loots and current action hotbar items should be added to the current memory cell state (backpack). The candidate $\tilde{C_t}$ is the combination of the new loots and the current action hotbar.
$$ I_t = \sigma(X_t W_{xi} + H_{t-1} W_{hi} + b_i) $$ $$ \tilde{C_t} = tanh(X_t W_{xc} + H_{t-1} W_{hc} + b_c) $$Output gate $O_t$
Now that you defeated the volcano dungeon boss, you are ready to move to the new dungeon, a sea (water) dungeon. Before you go to the new dungeon, you want to make sure that you are ready. So you look at your backpack (latest state), and your current action hotbar, then update the items in your current action hotbar. Since you are now moving from a volcano dungeon to a sea dungeon, you will replace your water elemental converter to a lightning elemental converter (lightning is effective against water).
The output gate determines which specific parts of the memory cell state should influence the output (hidden state) at the current time step.
$$ O_t = \sigma(X_t W_{xo} + H_{t-1} W_{ho} + b_o) $$Long Term and Short Term Memory
Now that we have gone over the three gates above, let’s talk about where LSTM got its name.
Cell State $C_t$ (Long Term Memory)
The cell state acts as the long-term memory. In our sentence analogy at the very start, the cell state carries the context of the word “Knight” so that our model can predict what weapon our subject wields when asked in the later part story. In our RPG analogy, it’s the backpack that allows our character to store and bring items we got from level 1 all the way to level 99.
How the Gates Regulate the Cell State
The cell state is controlled and protected by the three gates.
- Forget gate – decides what old information to throw away from the previous cell state $C_{t-1}$.
- Input gate – decides what new information from the current input and hidden state should be added to the cell state.
- Output gate – determines what part of the updated cell state becomes the visible hidden state $h_t$.
Cell State Update
The new cell state $C_t$ is calculated by combining the filtered old state and the new candidate state:
$$ C_t = F_t \odot C_{t-1} + I_t \odot \tilde{C_t} $$Hidden State $H_t$ (Short Term Memory)
The hidden state acts as the short-term memory. In our RPG analogy, it’s the Action Hotbar where we put items that our character needs immediately (like healing potions).
$$ H_t = O_t \odot \text{tanh}(C_t) $$Remember that the output gate’s value $O_t$ is in the range of (0,1). Whenever the output gate is close to 1, we allow the cell state to impact the rest of the neural network. Whereas for output gate values close to 0, we prevent the cell state from impacting other layers of the network at the current time step.
A crucial takeaway to remember is that an output gate of 0, $O_t = 0$, does not mean that the information is thrown away (that’s the Forget Gate’s $F_t$ job). It simply means that the information remains in the long term memory (cell state) $C_t$, hidden from the current time step $H_t$ and waiting for its turn to be helpful.
How LSTM addresses the Vanishing Gradient Problem
At the start of this blog post, I showed why the standard RNN fails using the BPTT multiplication table (0.5 x 0.5 x 0.5 x …). I just want to explicitly show how LSTM addresses this.
Let’s look again on how we compute for the long-term memory (cell state).
$$ C_t = F_t \odot C_{t-1} + I_t \odot \tilde{C_t} $$We see that the old long-term memory $C_{t-1}$ is carried forward via addition rather than continuous matrix multiplication. So the error gradients can flow backward intact without vanishing to 0.0078 (0.78%)!