Skip to content

batch_normalisation.md

Latest commit

 

History

History
299 lines (223 loc) · 10.4 KB

batch_normalisation.md

File metadata and controls

299 lines (223 loc) · 10.4 KB
Raw

A Beginner's Guide to Batch Normalization: Making Deep Learning Training Smooth and Fast

If you're learning deep learning, you've probably noticed that training neural networks can be frustratingly slow and unpredictable. Sometimes your model learns quickly, sometimes it gets stuck, and sometimes it seems to "forget" what it learned. Batch Normalization is one of the most important techniques that can solve these problems and make your training much more reliable.

This guide will walk you through what Batch Normalization is, why it works so well, and most importantly, how to use it effectively in TensorFlow.

The Problem: Why Neural Networks Are Hard to Train

The Internal Covariate Shift Problem

Imagine you're learning to catch baseballs. You practice for weeks catching balls thrown at chest height. Then suddenly, someone starts throwing balls at different heights - some high, some low, some to your left, some to your right. You'd have to relearn how to position yourself for each new type of throw.

Neural networks face a similar problem. Here's what happens:

  1. Layer 1 learns to recognize certain patterns in your data
  2. Layer 2 learns to work with the outputs from Layer 1
  3. But as Layer 1's weights change during training, the outputs it produces change too
  4. Now Layer 2 has to constantly readjust to these changing inputs
  5. This creates a cascade effect through the entire network

This phenomenon is called internal covariate shift - the distribution of inputs to each layer keeps shifting as the network trains.

Why This Causes Problems

This shifting creates several issues:

Slow Training: Each layer wastes time readjusting to the changing inputs instead of learning the actual task.

Vanishing/Exploding Gradients: As signals pass through many layers, they can become too small (vanishing) or too large (exploding). The shifting distributions make this worse.

Sensitive to Initialization: Small changes in how you initialize weights can make the difference between a model that learns and one that doesn't.

The Solution: Batch Normalization

The Core Idea

Batch Normalization solves this by ensuring that the inputs to each layer have a consistent, standardized distribution. Think of it as giving each layer a "stable foundation" to learn from.

Here's the key insight: instead of letting each layer deal with wildly varying inputs, we normalize them to have:

  • Mean = 0 (centered around zero)
  • Standard deviation = 1 (consistent spread)

But here's the clever part: we also give the network the ability to "undo" this normalization if needed, using learnable parameters.

The Math (Don't Worry, It's Simple!)

For each mini-batch of data, here's what happens:

Step 1: Calculate batch statistics

μ = (1/m) × Σ(xᵢ)     // Mean of the batch
σ² = (1/m) × Σ(xᵢ - μ)²  // Variance of the batch

Step 2: Normalize

x̂ᵢ = (xᵢ - μ) / √(σ² + ε)

This gives us normalized values with mean=0 and std=1. The tiny ε (epsilon) prevents division by zero.

Step 3: Scale and shift (the learnable part)

yᵢ = γ × x̂ᵢ + β

Here, γ (gamma) and β (beta) are learnable parameters. The network can learn to set these to whatever values work best.

Why the Scale and Shift Parameters Matter

You might wonder: "Why normalize just to potentially un-normalize?"

The answer is flexibility. Sometimes the network works best with normalized inputs (γ=1, β=0). But sometimes it needs a different mean and standard deviation. By making γ and β learnable, we let the network decide what distribution works best for each layer.

Practical Implementation in TensorFlow

Basic Usage Pattern

The most common pattern is to apply Batch Normalization before the activation function:

import tensorflow as tf

# Basic pattern: Dense -> BatchNorm -> Activation
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, input_shape=(784,)),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Activation('relu'),
    
    tf.keras.layers.Dense(32),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Activation('relu'),
    
    tf.keras.layers.Dense(10, activation='softmax')  # Output layer
])

Complete Example: MNIST with and without Batch Normalization

Let's see the difference in action:

import tensorflow as tf
import matplotlib.pyplot as plt

# Load and preprocess MNIST data
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train.reshape(-1, 784).astype('float32') / 255.0
x_test = x_test.reshape(-1, 784).astype('float32') / 255.0

# Model WITHOUT Batch Normalization
model_no_bn = tf.keras.Sequential([
    tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(32, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

model_no_bn.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

print("Training WITHOUT Batch Normalization:")
history_no_bn = model_no_bn.fit(
    x_train, y_train,
    epochs=10,
    batch_size=32,
    validation_data=(x_test, y_test),
    verbose=1
)

# Model WITH Batch Normalization
model_with_bn = tf.keras.Sequential([
    tf.keras.layers.Dense(128, input_shape=(784,)),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Activation('relu'),
    
    tf.keras.layers.Dense(64),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Activation('relu'),
    
    tf.keras.layers.Dense(32),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Activation('relu'),
    
    tf.keras.layers.Dense(10, activation='softmax')
])

model_with_bn.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

print("\nTraining WITH Batch Normalization:")
history_with_bn = model_with_bn.fit(
    x_train, y_train,
    epochs=10,
    batch_size=32,
    validation_data=(x_test, y_test),
    verbose=1
)

# Compare the results
plt.figure(figsize=(12, 4))

plt.subplot(1, 2, 1)
plt.plot(history_no_bn.history['accuracy'], label='No BN - Training')
plt.plot(history_no_bn.history['val_accuracy'], label='No BN - Validation')
plt.plot(history_with_bn.history['accuracy'], label='With BN - Training')
plt.plot(history_with_bn.history['val_accuracy'], label='With BN - Validation')
plt.title('Accuracy Comparison')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(history_no_bn.history['loss'], label='No BN - Training')
plt.plot(history_no_bn.history['val_loss'], label='No BN - Validation')
plt.plot(history_with_bn.history['loss'], label='With BN - Training')
plt.plot(history_with_bn.history['val_loss'], label='With BN - Validation')
plt.title('Loss Comparison')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()

plt.tight_layout()
plt.show()

For Convolutional Networks

Batch Normalization is especially powerful in CNNs:

# CNN with Batch Normalization
cnn_model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, (3, 3), input_shape=(28, 28, 1)),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Activation('relu'),
    tf.keras.layers.MaxPooling2D((2, 2)),
    
    tf.keras.layers.Conv2D(64, (3, 3)),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Activation('relu'),
    tf.keras.layers.MaxPooling2D((2, 2)),
    
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(64),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Activation('relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

When and Where to Use Batch Normalization

Almost Always Use It When:

  • Training deep networks (more than 3-4 layers)
  • Using activation functions like ReLU, Leaky ReLU, or ELU
  • Working with image data (CNNs)
  • Training is slow or unstable

Placement Guidelines:

  1. Most common: Dense/Conv → BatchNorm → Activation
  2. Alternative: Dense/Conv → Activation → BatchNorm
  3. Don't use on output layers (usually)

When to Be Careful:

  • Very small batch sizes (< 8): BatchNorm becomes less effective
  • Recurrent networks: Use LayerNormalization instead
  • When using Dropout: Place BatchNorm before Dropout

Key Benefits You'll See

1. Faster Training

You can often use higher learning rates, which means faster convergence.

2. Less Sensitive to Initialization

Random weight initialization becomes less critical.

3. Built-in Regularization

The noise from batch statistics provides a mild regularization effect.

4. More Stable Training

Your loss curves will be smoother and more predictable.

Common Pitfalls and Solutions

Pitfall 1: Forgetting Training vs. Inference Mode

# During training: uses batch statistics
# During inference: uses moving averages

# TensorFlow handles this automatically, but be aware:
model.fit(...)  # Training mode
model.predict(...)  # Inference mode

Pitfall 2: Wrong Placement

# Wrong: BatchNorm after activation might reduce effectiveness
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.BatchNormalization(),

# Better: BatchNorm before activation
tf.keras.layers.Dense(64),
tf.keras.layers.BatchNormalization(),
tf.keras.layers.Activation('relu'),

Pitfall 3: Small Batch Sizes

If you must use small batches, consider:

  • Using Layer Normalization instead
  • Accumulating gradients over multiple mini-batches
  • Using Group Normalization

Quick Reference: When to Use What

Scenario Recommendation
Deep feedforward networks Use BatchNorm before activation
CNNs Use BatchNorm after Conv layers
RNNs/LSTMs Use LayerNormalization instead
Very small batches Consider LayerNormalization
Transfer learning Keep BatchNorm layers frozen initially

Conclusion

Batch Normalization is one of the most impactful techniques in modern deep learning. It makes training more stable, faster, and less dependent on perfect hyperparameter tuning.

Key takeaways:

  1. Use it in almost all deep networks
  2. Place it before activation functions
  3. It normalizes inputs to each layer
  4. It includes learnable parameters (γ, β)
  5. TensorFlow handles the training/inference differences automatically

Start adding tf.keras.layers.BatchNormalization() to your models and you'll likely see immediate improvements in training speed and stability. It's one of those techniques that "just works" and makes your deep learning journey much smoother.