TensorFlow LSTM Layer
Long Short-Term Memory (LSTM) is an improved recurrent unit that solves the vanishing gradient problem of SimpleRNN. It introduces a separate cell state — a "memory highway" that runs alongside the hidden state — and three gating mechanisms that control what information gets added, removed, or passed forward. LSTMs remember relevant information across hundreds of time steps and remain the gold standard for many sequence tasks.
The Filing Cabinet Analogy
Imagine a researcher reading a long document. They maintain a notebook (cell state) where they jot key facts. When they read a new paragraph, they decide: which old notes are no longer relevant (forget gate), which new facts from this paragraph are worth adding (input gate), and which notes should inform their current thinking (output gate). LSTMs do exactly this for sequences.
The Three Gates
At each time step, LSTM computes with these gates:
FORGET GATE — What to erase from cell state
f_t = sigmoid(W_f × [h_{t-1}, x_t] + b_f)
Values near 0: forget this information
Values near 1: keep this information
INPUT GATE — What new information to store
i_t = sigmoid(W_i × [h_{t-1}, x_t] + b_i)
g_t = tanh(W_g × [h_{t-1}, x_t] + b_g)
New candidate values to potentially add: g_t
How much to add: i_t × g_t
OUTPUT GATE — What to expose as hidden state
o_t = sigmoid(W_o × [h_{t-1}, x_t] + b_o)
h_t = o_t × tanh(c_t)
CELL STATE UPDATE:
c_t = f_t × c_{t-1} + i_t × g_t
↑ ↑
(keep old) (add new)
Visualizing the Cell State Highway
Cell State (long-term memory):
c_0 ──[× f_1]──[+ i_1×g_1]──► c_1 ──[× f_2]──[+ i_2×g_2]──► c_2 ...
Hidden State (short-term memory exposed at each step):
h_0 ──► [LSTM cell] ──► h_1 ──► [LSTM cell] ──► h_2 ...
▲ ▲
│ x_1 (input) │ x_2 (input)
The cell state c_t flows through with minimal transformation.
This unobstructed highway lets gradients flow back
through many time steps without vanishing.
LSTM in TensorFlow
import tensorflow as tf
model = tf.keras.Sequential([
tf.keras.layers.Embedding(input_dim=10000, output_dim=128, input_length=200),
tf.keras.layers.LSTM(units=256),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dropout(0.3),
tf.keras.layers.Dense(1, activation='sigmoid')
])
model.compile(optimizer='adam',
loss='binary_crossentropy',
metrics=['accuracy'])
model.summary()
Key LSTM Parameters
tf.keras.layers.LSTM(
units=256, # Size of hidden state (and cell state)
activation='tanh', # Activation for cell state update
recurrent_activation='sigmoid', # Activation for gates
return_sequences=False, # False: return final h_t only
# True: return h_t at every time step
return_state=False, # True: also return h_t and c_t separately
dropout=0.0, # Dropout on input connections
recurrent_dropout=0.0, # Dropout on recurrent connections
stateful=False # True: carry state across batches
)
Stacking LSTM Layers
model = tf.keras.Sequential([
tf.keras.layers.Embedding(10000, 128, input_length=200),
# First LSTM — return sequences for next LSTM to consume
tf.keras.layers.LSTM(256, return_sequences=True, dropout=0.2,
recurrent_dropout=0.1),
# Second LSTM — return final state only
tf.keras.layers.LSTM(128, dropout=0.2),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(5, activation='softmax') # 5-class classification
])
Bidirectional LSTM
# Reads sequence both forward and backward
# Doubles the output size
model = tf.keras.Sequential([
tf.keras.layers.Embedding(10000, 128, input_length=200),
tf.keras.layers.Bidirectional(
tf.keras.layers.LSTM(128, return_sequences=True)
),
# output: (batch, 200, 256) ← 128 forward + 128 backward
tf.keras.layers.Bidirectional(
tf.keras.layers.LSTM(64)
),
# output: (batch, 128)
tf.keras.layers.Dense(1, activation='sigmoid')
])
Returning Cell State and Hidden State
# When return_state=True, LSTM returns 3 tensors: # output, final_hidden_state, final_cell_state lstm_layer = tf.keras.layers.LSTM(256, return_state=True) inputs = tf.keras.Input(shape=(100, 64)) output, h_state, c_state = lstm_layer(inputs) print(output.shape) # (batch, 256) — same as h_state here print(h_state.shape) # (batch, 256) — final hidden state print(c_state.shape) # (batch, 256) — final cell state # Use h_state and c_state to initialize a decoder (encoder-decoder models)
Complete Example: Movie Review Sentiment
import tensorflow as tf
import numpy as np
# Load IMDB dataset (25,000 movie reviews, binary sentiment)
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.imdb.load_data(
num_words=10000
)
# Pad sequences to same length
x_train = tf.keras.preprocessing.sequence.pad_sequences(
x_train, maxlen=200, padding='post', truncating='post'
)
x_test = tf.keras.preprocessing.sequence.pad_sequences(
x_test, maxlen=200, padding='post', truncating='post'
)
# Build LSTM model
model = tf.keras.Sequential([
tf.keras.layers.Embedding(10000, 128, input_length=200),
tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(64)),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dropout(0.4),
tf.keras.layers.Dense(1, activation='sigmoid')
])
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
model.fit(x_train, y_train,
epochs=5,
batch_size=64,
validation_split=0.2)
loss, acc = model.evaluate(x_test, y_test)
print(f"Test accuracy: {acc:.2%}")
# Typically achieves ~87–89% accuracy
LSTM vs SimpleRNN vs GRU
Property SimpleRNN LSTM GRU ──────────────────────────────────────────────────────────── Gates 0 3 (f, i, o) 2 (r, z) Memory Hidden state Cell + hidden Hidden only Parameters (units=u) u² 4u² 3u² Long-range memory Poor Excellent Good Training speed Fastest Slowest Fast Accuracy (long seq) Low High High (≈ LSTM) Preferred use Short seqs All seq tasks When speed matters ────────────────────────────────────────────────────────────
When LSTM Struggles
- Very long sequences (>1000 steps) — even LSTM degrades; use Transformers
- Parallel training — LSTM is sequential by nature; Transformers run in parallel
- Small datasets — fewer parameters (GRU) often outperforms LSTM
LSTM is the workhorse of sequence modeling and one of the most proven architectures in deep learning history. Its gating mechanism solves the core weakness of SimpleRNN, making it effective for text classification, sentiment analysis, time series forecasting, and language modeling. The next topic covers the GRU layer — a streamlined version of LSTM that achieves comparable accuracy with fewer parameters and faster training.