SimpleRNN, LSTM, and GRU

The recurrent cell family

RNNs process sequences one step at a time, maintaining a hidden state that summarizes everything seen so far. Each step takes the current input and the previous hidden state, produces a new hidden state, and (optionally) an output.

SimpleRNN is the most basic cell. It struggles past 20–30 time steps because gradients vanish or explode through deep unrolled computation.

LSTM (Long Short-Term Memory) solves the vanishing gradient problem with a two-stream architecture: a cell state (long-term memory) and a hidden state (working memory). Three gates — forget, input, output — control information flow. LSTMs reliably model dependencies of hundreds of time steps.

GRU (Gated Recurrent Unit) is a streamlined LSTM with two gates instead of three, merging cell and hidden state. ~25% fewer parameters than equivalent LSTM, similar accuracy on most tasks, faster training. Set reset_after=True for the CuDNN-optimized GPU implementation — up to 5× faster than the default.

Code

Three RNN flavors·python

import tensorflow as tf
from tensorflow.keras import layers

# SimpleRNN — for short sequences only
rnn = layers.SimpleRNN(units=64, return_sequences=False, activation='tanh')

# LSTM — long-range memory via cell state + 3 gates
lstm = layers.LSTM(
    units=128,
    return_sequences=True,    # True for stacking, False for final-only
    return_state=False,       # True returns (output, hidden_state, cell_state)
    dropout=0.2,
    recurrent_dropout=0.1,
)

# GRU — fewer parameters, similar accuracy, CuDNN-friendly
gru = layers.GRU(
    units=128,
    return_sequences=True,
    reset_after=True,         # ALWAYS True for CuDNN GPU acceleration
)

# Parameter counts for units=128, input_size=128:
# SimpleRNN: ~32K
# GRU:       ~99K
# LSTM:      ~132K

Stacked LSTM for sentiment·python

import tensorflow as tf
from tensorflow.keras import layers

model = tf.keras.Sequential([
    layers.TextVectorization(max_tokens=10000, output_sequence_length=200),
    layers.Embedding(10000 + 2, 128, mask_zero=True),
    layers.LSTM(64, return_sequences=True),    # stacking → return_sequences=True
    layers.LSTM(32),                           # final → return last state only
    layers.Dense(64, activation='relu'),
    layers.Dropout(0.5),
    layers.Dense(1, activation='sigmoid'),
])

model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

The recurrent cell family

Code

Progress

Comments 0