Batch Shuffling & Mini-Batch Generator

Data Processing

Medium

Problem

Theory

In machine learning, training on entire datasets at once is often impractical due to memory constraints. Mini-batch training processes data in small, manageable chunks called batches. A batch generator creates these batches on-the-fly, enabling efficient training of large-scale models.

Why Batch Training Matters

When training neural networks, the choice of batch size affects three critical factors:

• Memory Efficiency: Loading 1 million samples into GPU memory simultaneously is impossible. Batches of 32-256 samples fit comfortably within typical hardware limits.

• Gradient Quality: The gradient computed from a batch approximates the true gradient over the full dataset. The approximation quality depends on batch size and data diversity.

\nabla_\theta L_{batch} \approx \nabla_\theta L_{full}

• Training Dynamics: Smaller batches introduce noise that can help escape local minima, while larger batches provide more stable gradient estimates but may converge to sharper minima.

• Computational Throughput: GPUs are optimized for parallel operations on fixed-size tensors. Batching allows efficient vectorized computation.

The Mathematics of Batching

Given a dataset of N samples and batch size B:

• Number of complete batches: \lfloor N / B \rfloor

• Remaining samples: N \mod B

• Total batches including partial: \lceil N / B \rceil

The drop_last parameter determines whether to discard the final incomplete batch:

• drop_last=True: Exactly \lfloor N / B \rfloor batches, each with B samples
• drop_last=False: May have one smaller batch at the end

Shuffling and Randomization

Shuffling data between epochs prevents the model from learning order-dependent patterns:

• Why shuffle? If data is ordered (e.g., all class A samples first, then class B), the model sees biased batches and learns spurious correlations.

• Index Shuffling: Create a random permutation of indices \pi = [\pi_1, \pi_2, ..., \pi_N] and access data via X[\pi]. Memory efficient since only indices are stored and shuffled.

• Full Shuffle: Physically reorder the data array. More expensive but may improve cache locality during iteration.

• Fisher-Yates Algorithm: The standard algorithm for generating uniform random permutations in O(N) time by iteratively swapping each position with a random earlier position.

Handling Paired Data

When features X and labels y must stay aligned, shuffling must maintain correspondence:

• Correct approach: Shuffle indices once, then apply the same index permutation to both X and y arrays. This preserves the pairing between features and labels.

• Incorrect approach: Shuffling X and y independently destroys the pairing and results in random label assignments.

• Practical consideration: Store paired data together or use a single index array that references both.

Epoch vs Iteration Terminology

• Iteration: Processing one batch through the network (one forward pass + one backward pass + one parameter update)

• Epoch: Processing the entire dataset once

• Relationship: For N=10000 samples and B=100, one epoch consists of 100 iterations

• Training duration: Often specified in epochs (e.g., "train for 50 epochs") which translates to epochs × iterations_per_epoch total updates

Stratified Batching

For imbalanced classification, random batching may produce batches with skewed class distributions:

• Problem: If dataset has 99% class A and 1% class B, some batches may contain zero class B samples, providing no learning signal for the minority class.

• Solution: Stratified batching ensures each batch approximately reflects the overall class distribution.

• Implementation idea: Sample proportionally from each class when constructing batches, or oversample minority classes.

Generator Pattern Benefits

A generator produces batches lazily rather than creating all batches upfront:

• Memory efficiency: Only the current batch occupies memory, not the entire batched dataset

• Infinite iteration: Generators can cycle through data indefinitely for training

• On-the-fly processing: Data augmentation, normalization, or other transforms can be applied per-batch

• Large dataset handling: Enables training on datasets larger than available RAM by loading batches from disk

Practical Considerations

• Batch size selection: Powers of 2 (32, 64, 128, 256) often perform best due to GPU memory alignment

• Reproducibility: Setting random seeds before shuffling ensures reproducible batch compositions

• Last batch handling: Incomplete final batches can cause issues with batch normalization layers that require consistent batch sizes

• Multi-epoch iteration: Re-shuffle at the start of each epoch to ensure different batch compositions

Where Batch Generators Show Up

• Deep Learning Frameworks: PyTorch DataLoader, TensorFlow tf.data.Dataset, Keras Sequence all implement batch generation with prefetching and parallel loading

• Large-Scale Training: Training GPT-style models on terabytes of text requires streaming batches from distributed storage

• Online Learning: Systems processing continuous data streams accumulate samples into batches before model updates

• Data Augmentation Pipelines: Batch generators apply random transformations (rotation, flipping, noise) to each batch, effectively creating infinite training data

• Distributed Training: Multi-GPU setups require batch generators that partition data across workers without sample duplication within an epoch

• Transfer Learning: Fine-tuning pretrained models on new datasets using batched gradient descent

• Hyperparameter Search: Different batch sizes tested during hyperparameter optimization to find optimal training configuration

Examples

Hints

1. Shuffle np.arange(len(X)) with np.random.default_rng(seed).
2. Slice the shuffled indices inside range(0, len(indices), batch_size).

Requirements

• Return a Python generator
• Apply one seeded permutation to both X and y
• Yield tuples of NumPy arrays
• Respect drop_last without modifying the inputs

Constraints

• X and y contain the same positive number of samples
• batch_size is a positive integer
• Use NumPy only

Starter Code

import numpy as np

def batch_generator(X: list, y: list, batch_size: int, seed: int = 42, drop_last: bool = False):
    """
    Returns a generator of (X_batch, y_batch) tuples.
    """
    # Write code here
    pass

Test Cases