Introduction

Batch normalization is a technique that normalizes the inputs to each layer, stabilizing training and often allowing higher learning rates. Applying it to LSTM networks is not as straightforward as inserting a layer into a feedforward model because LNNs process sequential data across time steps. This article covers practical approaches to adding normalization between LSTM layers in TensorFlow/Keras, explores Layer Normalization as a recurrence-friendly alternative, and demonstrates a custom LSTM cell with built-in normalization.

Batch Normalization Between LSTM Layers

The simplest approach is to place a BatchNormalization layer between stacked LSTM layers. When you stack multiple LSTM layers, you can normalize the output of one LSTM before feeding it into the next. The key requirement is that the first LSTM must return full sequences so the next LSTM receives input at every time step.

1import tensorflow as tf
2from tensorflow.keras import layers, models
3
4model = models.Sequential([
5    layers.LSTM(128, return_sequences=True, input_shape=(50, 10)),
6    layers.BatchNormalization(),
7    layers.LSTM(64, return_sequences=False),
8    layers.BatchNormalization(),
9    layers.Dense(32, activation='relu'),
10    layers.Dense(1)
11])
12
13model.compile(optimizer='adam', loss='mse')
14model.summary()

The BatchNormalization layer computes statistics across the batch dimension. During training it uses the current mini-batch mean and variance, while during inference it uses running averages accumulated during training.

Why Layer Normalization Is Often Better for RNNs

Batch normalization computes statistics across samples in a mini-batch. In recurrent networks, this creates problems because batch statistics can vary significantly across time steps, and small batch sizes lead to noisy estimates. Layer Normalization computes statistics across the feature dimension for each individual sample, making it independent of batch size and more stable for sequential models.

1model = models.Sequential([
2    layers.LSTM(128, return_sequences=True, input_shape=(50, 10)),
3    layers.LayerNormalization(),
4    layers.LSTM(64, return_sequences=False),
5    layers.LayerNormalization(),
6    layers.Dense(32, activation='relu'),
7    layers.Dense(1)
8])
9
10model.compile(optimizer='adam', loss='mse')

Layer Normalization works well with variable-length sequences and batch size of 1, which are common scenarios in RNN applications such as text generation and time-series forecasting.

Custom LSTM Cell with Normalization

For finer control, you can create a custom LSTM cell that applies normalization inside the recurrence, normalizing the gate activations at every time step rather than only between layers.

1class BNLSTMCell(layers.AbstractRNNCell):
2    def __init__(self, units, **kwargs):
3        super().__init__(**kwargs)
4        self.units = units
5        self.dense_i = layers.Dense(4 * units, use_bias=False)
6        self.dense_h = layers.Dense(4 * units, use_bias=False)
7        self.bn_i = layers.BatchNormalization()
8        self.bn_h = layers.BatchNormalization()
9        self.bn_c = layers.BatchNormalization()
10
11    @property
12    def state_size(self):
13        return [self.units, self.units]
14
15    @property
16    def output_size(self):
17        return self.units
18
19    def call(self, inputs, states, training=None):
20        h_prev, c_prev = states
21        x_transformed = self.bn_i(self.dense_i(inputs), training=training)
22        h_transformed = self.bn_h(self.dense_h(h_prev), training=training)
23        gates = x_transformed + h_transformed
24
25        i, f, o, g = tf.split(gates, 4, axis=-1)
26        i = tf.sigmoid(i)
27        f = tf.sigmoid(f)
28        o = tf.sigmoid(o)
29        g = tf.tanh(g)
30
31        c_new = f * c_prev + i * g
32        c_normed = self.bn_c(c_new, training=training)
33        h_new = o * tf.tanh(c_normed)
34        return h_new, [h_new, c_new]

You can then use this cell with the RNN wrapper layer.

1cell = BNLSTMCell(128)
2model = models.Sequential([
3    layers.RNN(cell, return_sequences=False, input_shape=(50, 10)),
4    layers.Dense(32, activation='relu'),
5    layers.Dense(1)
6])
7
8model.compile(optimizer='adam', loss='mse')

Training Example

Here is a complete training example that uses batch normalization between LSTM layers on synthetic time-series data.

1import numpy as np
2
3# Generate synthetic sequential data
4x_train = np.random.randn(1000, 50, 10).astype(np.float32)
5y_train = np.random.randn(1000, 1).astype(np.float32)
6x_val = np.random.randn(200, 50, 10).astype(np.float32)
7y_val = np.random.randn(200, 1).astype(np.float32)
8
9model = models.Sequential([
10    layers.LSTM(128, return_sequences=True, input_shape=(50, 10)),
11    layers.BatchNormalization(),
12    layers.Dropout(0.3),
13    layers.LSTM(64),
14    layers.BatchNormalization(),
15    layers.Dropout(0.3),
16    layers.Dense(32, activation='relu'),
17    layers.Dense(1)
18])
19
20model.compile(
21    optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
22    loss='mse'
23)
24
25history = model.fit(
26    x_train, y_train,
27    validation_data=(x_val, y_val),
28    epochs=20,
29    batch_size=64
30)

Combine BatchNormalization with Dropout for regularization. Place batch normalization before dropout so that the normalization statistics are computed on the full output before units are randomly zeroed out.

Common Pitfalls

• Using small batch sizes with BatchNormalization: Batch norm estimates become noisy with very small batches, leading to unstable training. Use Layer Normalization instead when batch sizes are below 16.
• Forgetting return_sequences on intermediate LSTMs: Stacking LSTM layers requires return_sequences=True on all layers except the last, otherwise the subsequent LSTM receives a single vector instead of a sequence.
• Ignoring training vs inference mode: BatchNormalization behaves differently during training and inference. Always pass training=True or training=False correctly when using custom training loops.
• Applying BN after activation in LSTM gates: Inside a custom cell, normalize before the activation functions, not after. Normalizing after sigmoid or tanh compresses the distribution and reduces the normalization benefit.
• Not freezing BN layers during fine-tuning: When fine-tuning a pretrained model with very few samples, freeze BatchNormalization layers by setting layer.trainable = False to prevent running statistics from being corrupted by the small fine-tuning dataset.

Summary

• Place BatchNormalization between stacked LSTM layers for a quick improvement in training stability.
• Prefer LayerNormalization over batch normalization for recurrent networks, especially with small or variable batch sizes.
• Build a custom LSTM cell with normalization inside the recurrence for gate-level control.
• Always set return_sequences=True on LSTM layers that feed into another recurrent or normalization layer.
• Combine normalization with dropout for regularization, placing normalization before dropout in the layer order.