The LSTM Cell Architecture

As we discussed previously, simple Recurrent Neural Networks face significant challenges when learning dependencies across long sequences, primarily due to vanishing and exploding gradients. The network's ability to carry relevant information from earlier time steps gets compromised. Long Short-Term Memory (LSTM) networks were specifically developed to address this limitation by introducing a more complex cell structure capable of maintaining a memory over extended periods.

At the core of an LSTM network is the LSTM cell. It replaces the simple transformation found in a standard RNN cell with a sophisticated system of gates and a dedicated cell state. This architecture allows the network to selectively add, remove, or retain information over time.

The LSTM cell has two primary components that enable this controlled information flow:

1. The Cell State ($C_t$): This is the central feature of the LSTM. You can think of it as an internal memory track running horizontally through the sequence of cells. Information can flow along this state line with only minor linear interactions. This structure makes it much easier for information (and gradients during training) to persist across many time steps without significant degradation.
2. Gates: These are neural network layers (typically using sigmoid $\sigma$ activation functions) that regulate the flow of information into and out of the cell state. Because the sigmoid function outputs values between 0 and 1, these gates act like filters: a value close to 0 means "let very little information pass," while a value close to 1 means "let most information pass." An LSTM cell contains three main gates:
   • Forget Gate: Decides which information to discard from the previous cell state ($C_{t-1}$).
   • Input Gate: Determines which new information from the current input ($x_t$) and previous hidden state ($h_{t-1}$) should be stored in the cell state.
   • Output Gate: Controls which parts of the current cell state ($C_t$) should be passed on as the hidden state ($h_t$) for the current time step.

Let's visualize how these components interact within a single LSTM cell at time step $t$:

Data flow and components within a single LSTM cell at time step $t$. It receives the current input $x_t$, the previous hidden state $h_{t-1}$, and the previous cell state $C_{t-1}$. It computes the new cell state $C_t$ and the new hidden state $h_t$. Circles represent operations (sigmoid $\sigma$, hyperbolic tangent $\tanh$, element-wise multiplication $\times$, element-wise addition $+$).

Information processing within the LSTM cell proceeds as follows:

1. Forget Gate ($f_t$): The cell first decides what information to throw away from the previous cell state, $C_{t-1}$. It looks at the previous hidden state $h_{t-1}$ and the current input $x_t$. These are passed through a sigmoid function $\sigma$. The output $f_t$ contains values between 0 and 1 for each number in $C_{t-1}$. A 1 represents "completely keep this," while a 0 represents "completely get rid of this." $f_t = \sigma(W_f [h_{t-1}, x_t] + b_f)$

2. Input Gate ($i_t$) and Candidate Values ($\tilde{C}_t$): Next, the cell decides what new information to store in the cell state. This involves two steps:

   • A sigmoid layer, called the "input gate layer," takes $h_{t-1}$ and $x_t$ and decides which values to update ($i_t$).
   • A $\tanh$ layer takes $h_{t-1}$ and $x_t$ and creates a vector of new candidate values, $\tilde{C}_t$, that could potentially be added to the state. $i_t = \sigma(W_i [h_{t-1}, x_t] + b_i)$ $\tilde{C}_t = \tanh(W_C [h_{t-1}, x_t] + b_C)$

3. Cell State Update ($C_t$): Now, the old cell state $C_{t-1}$ is updated into the new cell state $C_t$. The previous state $C_{t-1}$ is multiplied element-wise ($\odot$) by the forget vector $f_t$ (forgetting the selected parts). Then, the result of the element-wise multiplication $i_t \odot \tilde{C}_t$ (the new information, scaled by how much we decided to update) is added. $C_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t$

4. Output Gate ($o_t$) and Hidden State ($h_t$): Finally, the cell determines the output, which is the hidden state $h_t$. This output is a filtered version of the cell state.

   • First, a sigmoid layer takes $h_{t-1}$ and $x_t$ to decide which parts of the cell state to output ($o_t$).
   • Then, the newly computed cell state $C_t$ is passed through $\tanh$ (squashing values to be between -1 and 1).
   • This $\tanh(C_t)$ is multiplied element-wise ($\odot$) by the output gate's activation $o_t$. This produces the hidden state $h_t$. $o_t = \sigma(W_o [h_{t-1}, x_t] + b_o)$ $h_t = o_t \odot \tanh(C_t)$

In the equations above, $W_f, W_i, W_C, W_o$ represent the weight matrices and $b_f, b_i, b_C, b_o$ are the bias vectors, which are learned during training. The notation $[h_{t-1}, x_t]$ often signifies the concatenation of the two vectors.

This gated architecture, especially the distinct cell state $C_t$ that undergoes only minor additions and multiplications controlled by gates, is the reason LSTMs are much better at capturing long-range dependencies than simple RNNs. Information can be preserved over many steps, and the gates learn to control what information is relevant, effectively mitigating the vanishing gradient problem by providing more direct pathways for gradient propagation.