TensorFlow Batch Normalization
Batch Normalization (BatchNorm) normalizes the output of each layer so that it has a mean near zero and standard deviation near one during training. This prevents the inputs to each layer from shifting wildly as training progresses — a problem called internal covariate shift. BatchNorm is one of the most impactful techniques in modern deep learning: it dramatically speeds up training, allows higher learning rates, and acts as a regularizer, reducing the need for Dropout in many architectures.
The Problem BatchNorm Solves
Imagine training a network and the first layer's weights change slightly. The distribution of its outputs shifts — outputs that were typically around 0 might now hover around 5. The next layer has to re-adapt to this new distribution. Then its weights shift, and the layer after that has to re-adapt. This cascading adaptation makes training slow and unstable. BatchNorm fixes each layer's output distribution so that downstream layers always receive inputs in a consistent range.
Without BatchNorm: Layer 1 output distribution changes every step Layer 2 must constantly re-learn to handle new distributions Layer 3 same problem → Slow convergence, need very small learning rates With BatchNorm: Layer 1 output → normalized → consistent distribution Layer 2 always sees similar inputs Layer 3 same benefit → Faster convergence, can use larger learning rates
How BatchNorm Works
For each mini-batch and each feature dimension: Step 1: Compute batch mean μ_B = (1/m) × Σ x_i Step 2: Compute batch variance σ²_B = (1/m) × Σ (x_i - μ_B)² Step 3: Normalize x̂_i = (x_i - μ_B) / sqrt(σ²_B + ε) (ε is a small number like 0.001 to avoid division by zero) Step 4: Scale and shift (learnable) y_i = γ × x̂_i + β (γ = scale parameter, β = shift parameter — both learned during training) The model learns what mean and variance are optimal for each layer, rather than being forced to always have mean=0, std=1.
BatchNorm in TensorFlow
import tensorflow as tf
# Place BatchNorm AFTER the linear transformation, BEFORE the activation
model = tf.keras.Sequential([
tf.keras.layers.Dense(256, use_bias=False, input_shape=(100,)),
tf.keras.layers.BatchNormalization(),
tf.keras.layers.Activation('relu'),
tf.keras.layers.Dense(128, use_bias=False),
tf.keras.layers.BatchNormalization(),
tf.keras.layers.Activation('relu'),
tf.keras.layers.Dense(10, activation='softmax')
])
Note the use_bias=False in Dense layers before BatchNorm. BatchNorm has its own shift parameter β that acts as a bias, making the Dense bias redundant. Removing it reduces parameters slightly.
BatchNorm Placement Options
Option 1 — After linear, before activation (standard):
Dense → BatchNorm → Activation
Most common and generally performs best
Option 2 — After activation (original paper):
Dense → Activation → BatchNorm
Less common today but still used in some architectures
In CNNs:
Conv2D(use_bias=False) → BatchNormalization() → Activation('relu')
BatchNorm in CNNs
def conv_bn_relu(filters, kernel_size=3):
return tf.keras.Sequential([
tf.keras.layers.Conv2D(
filters, kernel_size,
padding='same',
use_bias=False # BN handles the bias
),
tf.keras.layers.BatchNormalization(),
tf.keras.layers.ReLU()
])
model = tf.keras.Sequential([
conv_bn_relu(32, input_shape=(32, 32, 3)),
tf.keras.layers.MaxPooling2D(2),
conv_bn_relu(64),
tf.keras.layers.MaxPooling2D(2),
conv_bn_relu(128),
tf.keras.layers.GlobalAveragePooling2D(),
tf.keras.layers.Dense(10, activation='softmax')
])
Training vs. Inference Behavior
BatchNorm behaves differently during training and inference — this is crucial to understand.
During TRAINING (training=True):
→ Normalize using the CURRENT BATCH's mean and variance
→ Also update running_mean and running_variance using exponential moving average
During INFERENCE (training=False):
→ Normalize using the STORED running_mean and running_variance
(accumulated from all training batches)
→ Do NOT use the current test batch statistics
→ This makes predictions consistent and independent of batch size
Why it matters:
If you pass training=True during inference, a batch of 1 sample
would be "normalized" to have mean=0, variance=1 — destroying the signal.
Always use training=False at inference time.
# In custom training loops, pass training explicitly output_train = model(x, training=True) # BN uses batch stats output_infer = model(x, training=False) # BN uses running stats # model.fit() handles this automatically # model.predict() and model.evaluate() use training=False automatically
BatchNorm Parameters
tf.keras.layers.BatchNormalization(
axis=-1, # Which axis to normalize (default: -1 = last axis = channels)
momentum=0.99, # Running average momentum (higher = slower update)
epsilon=0.001, # Added to variance to prevent division by zero
center=True, # Add β (shift) parameter
scale=True # Add γ (scale) parameter
)
Trainable and Non-Trainable Parameters in BatchNorm
For BatchNormalization on a layer with 256 features: Trainable parameters: γ (scale): 256 values ← adjusted by gradient descent β (shift): 256 values ← adjusted by gradient descent Non-trainable parameters: running_mean: 256 values ← updated by exponential moving average running_variance: 256 values ← NOT updated by gradient descent Total: 256 × 4 = 1024 parameters (512 trainable + 512 non-trainable)
Benefits of BatchNorm
Benefit How Much ────────────────────────────────────────────────────────────── Training speed 2–5× faster convergence Learning rate tolerance Can use 5–10× larger learning rates Regularization effect Reduces overfitting (like light Dropout) Weight initialization sensitivity Much less sensitive to bad initialization Gradient flow Improved through deep networks ──────────────────────────────────────────────────────────────
When BatchNorm Struggles
Situation Problem Alternative ────────────────────────────────────────────────────────────────────── Very small batch size (1–4) Noisy batch statistics LayerNormalization RNNs and LSTMs Time step statistics vary LayerNormalization Online learning (1 sample/step) Batch stats undefined LayerNormalization Highly variable sequence lengths Inconsistent normalization LayerNormalization ──────────────────────────────────────────────────────────────────────
# LayerNormalization for RNNs and Transformers # Normalizes across features (not across batch samples) tf.keras.layers.LayerNormalization(axis=-1, epsilon=1e-6)
Batch normalization is present in nearly every modern CNN architecture — ResNet, EfficientNet, MobileNet, and Inception all rely on it. Adding BatchNorm after each Conv2D or Dense layer consistently improves training speed and final accuracy with minimal code change. The next topic covers the SavedModel format — TensorFlow's universal model storage format for production deployment.