LSTM Gate Equations › Implementing LSTM from Scratch
Language AI HandbookAll chapters LSTM Gate Equations › Implementing LSTM from ScratchOne section · Sections in this chapter11 / 17 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Now let's implement an LSTM cell from scratch using only NumPy. This exercise solidifies understanding and reveals the computational structure hidden behind framework abstractions.
Activation Functions
First, we need the sigmoid and tanh activation functions:
import numpy as np
def sigmoid(x):
"""Numerically stable sigmoid function."""
return np.where(x >= 0, 1 / (1 + np.exp(-x)), np.exp(x) / (1 + np.exp(x)))
def tanh(x):
"""Hyperbolic tangent function."""
return np.tanh(x)The sigmoid implementation handles numerical stability by using different formulas for positive and negative inputs, avoiding overflow from large exponentials.
LSTM Cell Class
Now we implement the LSTM cell with explicit weight matrices for each component:
class LSTMCell:
"""A single LSTM cell implemented from scratch."""
def __init__(self, input_dim, hidden_dim):
self.input_dim = input_dim
self.hidden_dim = hidden_dim
# Initialize weights using Xavier initialization
scale_x = np.sqrt(2.0 / (input_dim + hidden_dim))
scale_h = np.sqrt(2.0 / (hidden_dim + hidden_dim))
# Forget gate parameters
self.W_f = np.random.randn(hidden_dim, input_dim) * scale_x
self.U_f = np.random.randn(hidden_dim, hidden_dim) * scale_h
self.b_f = np.ones(
hidden_dim
) # Initialize to 1 for better gradient flow
# Input gate parameters
self.W_i = np.random.randn(hidden_dim, input_dim) * scale_x
self.U_i = np.random.randn(hidden_dim, hidden_dim) * scale_h
self.b_i = np.zeros(hidden_dim)
# Candidate cell state parameters
self.W_c = np.random.randn(hidden_dim, input_dim) * scale_x
self.U_c = np.random.randn(hidden_dim, hidden_dim) * scale_h
self.b_c = np.zeros(hidden_dim)
# Output gate parameters
self.W_o = np.random.randn(hidden_dim, input_dim) * scale_x
self.U_o = np.random.randn(hidden_dim, hidden_dim) * scale_h
self.b_o = np.zeros(hidden_dim)
def forward(self, x_t, h_prev, c_prev):
"""
Forward pass through the LSTM cell.
Args:
x_t: Input at time t, shape (input_dim,)
h_prev: Previous hidden state, shape (hidden_dim,)
c_prev: Previous cell state, shape (hidden_dim,)
Returns:
h_t: New hidden state, shape (hidden_dim,)
c_t: New cell state, shape (hidden_dim,)
"""
# Forget gate
f_t = sigmoid(self.W_f @ x_t + self.U_f @ h_prev + self.b_f)
# Input gate
i_t = sigmoid(self.W_i @ x_t + self.U_i @ h_prev + self.b_i)
# Candidate cell state
c_tilde = tanh(self.W_c @ x_t + self.U_c @ h_prev + self.b_c)
# New cell state
c_t = f_t * c_prev + i_t * c_tilde
# Output gate
o_t = sigmoid(self.W_o @ x_t + self.U_o @ h_prev + self.b_o)
# New hidden state
h_t = o_t * tanh(c_t)
return h_t, c_t
def count_parameters(self):
"""Count total learnable parameters."""
return 4 * self.hidden_dim * (self.input_dim + self.hidden_dim + 1)Let's verify our implementation produces outputs with the correct shapes:
# Create an LSTM cell
input_dim = 10
hidden_dim = 20
lstm_cell = LSTMCell(input_dim, hidden_dim)
# Create sample inputs
x_t = np.random.randn(input_dim)
h_prev = np.zeros(hidden_dim)
c_prev = np.zeros(hidden_dim)
# Forward pass
h_t, c_t = lstm_cell.forward(x_t, h_prev, c_prev)LSTM Cell Configuration: Input dimension: 10 Hidden dimension: 20 Total parameters: 2,480 Input shapes: x_t: (10,) h_prev: (20,) c_prev: (20,) Output shapes: h_t: (20,) c_t: (20,) Sample output values: h_t[:5]: [-0.0291864 0.00752907 -0.19155684 -0.12173972 -0.03181409] c_t[:5]: [-0.26307027 0.01370535 -0.27035174 -0.24262702 -0.10380498]
The hidden state values are bounded between and because they're computed as the product of a sigmoid (range ) and a tanh (range ). The cell state values can exceed this range since they accumulate through addition.
Processing a Sequence
An LSTM processes sequences by applying the cell repeatedly, passing the hidden and cell states from one step to the next:
class LSTM:
"""LSTM layer that processes sequences."""
def __init__(self, input_dim, hidden_dim):
self.cell = LSTMCell(input_dim, hidden_dim)
self.hidden_dim = hidden_dim
def forward(self, x_sequence):
"""
Process a sequence through the LSTM.
Args:
x_sequence: Input sequence, shape (seq_len, input_dim)
Returns:
hidden_states: All hidden states, shape (seq_len, hidden_dim)
final_cell_state: Final cell state, shape (hidden_dim,)
"""
seq_len = x_sequence.shape[0]
# Initialize states
h_t = np.zeros(self.hidden_dim)
c_t = np.zeros(self.hidden_dim)
# Store all hidden states
hidden_states = []
# Process each time step
for t in range(seq_len):
h_t, c_t = self.cell.forward(x_sequence[t], h_t, c_t)
hidden_states.append(h_t)
return np.array(hidden_states), c_tLet's process a sample sequence and examine how the hidden state evolves:
# Create LSTM and sample sequence
lstm = LSTM(input_dim=10, hidden_dim=20)
sequence = np.random.randn(15, 10) # 15 time steps, 10-dimensional input
# Process the sequence
hidden_states, final_cell = lstm.forward(sequence)Sequence length: 15 Input dimension: 10 Hidden dimension: 20 Hidden states shape: (15, 20) Final cell state shape: (20,)
The LSTM produces a hidden state at each time step. For sequence classification, you typically use the final hidden state. For sequence-to-sequence tasks, you use all hidden states.
Visualizing Gate Activations
Let's visualize how the gates respond to a sequence, which reveals how the LSTM decides what to remember and forget:
class LSTMWithGateHistory(LSTM):
"""LSTM that records gate activations for analysis."""
def forward_with_gates(self, x_sequence):
"""Process sequence and return gate activations."""
seq_len = x_sequence.shape[0]
h_t = np.zeros(self.hidden_dim)
c_t = np.zeros(self.hidden_dim)
# Store gate activations
forget_gates = []
input_gates = []
output_gates = []
cell_states = []
for t in range(seq_len):
x_t = x_sequence[t]
cell = self.cell
# Compute gates explicitly
f_t = sigmoid(cell.W_f @ x_t + cell.U_f @ h_t + cell.b_f)
i_t = sigmoid(cell.W_i @ x_t + cell.U_i @ h_t + cell.b_i)
c_tilde = tanh(cell.W_c @ x_t + cell.U_c @ h_t + cell.b_c)
c_t = f_t * c_t + i_t * c_tilde
o_t = sigmoid(cell.W_o @ x_t + cell.U_o @ h_t + cell.b_o)
h_t = o_t * tanh(c_t)
forget_gates.append(f_t.mean())
input_gates.append(i_t.mean())
output_gates.append(o_t.mean())
cell_states.append(np.abs(c_t).mean())
return {
"forget": np.array(forget_gates),
"input": np.array(input_gates),
"output": np.array(output_gates),
"cell_magnitude": np.array(cell_states),
}
The visualization reveals several key behaviors. The forget gate starts high (near 1) because we initialized its bias to 1, causing the LSTM to preserve information by default. The input and output gates show more variation as they learn to selectively write and read. The cell state magnitude tends to grow over time as information accumulates, though the gates prevent unbounded growth.