Batch Processing - From Neuron to GPT

In each iteration so far, we trained with a single sample. This approach is called stochastic gradient descent (SGD). It is simple, but it cannot make full use of NumPy’s vectorized, parallel computation. In fact, we can send several samples into the network model at once for training. This approach is called batch processing and is the standard practice in modern deep learning.

Depending on how many samples are used each time, gradient descent falls into three categories:

Stochastic gradient descent (SGD): each iteration uses one sample;
Batch gradient descent (Batch GD): each iteration uses all the samples. In practice, batch gradient descent is rarely used for training, because it needs a lot of memory and runs very slowly;
Mini-batch gradient descent (Mini-batch GD): each iteration uses a few samples. Mini-batch gradient descent makes full use of the vectorized, parallel computation of NumPy (and other computing libraries), and is the usual way to train models today.

Convergence¶

Convergence means that during model training, as the network model iterates, its loss no longer decreases significantly and the model parameters become stable.

Stochastic gradient descent: each iteration learns from only one sample, so the gradient fluctuates a lot, the direction of convergence tends to swing from side to side, and convergence is slow;
Batch gradient descent: each iteration learns from all the data, so it can keep a relatively correct direction of convergence, and convergence is fast;
Mini-batch gradient descent: a compromise that balances efficiency and stability.

Generalization¶

Generalization is an important measure of how well a model has been trained: the ability of the network model to handle new data. Put simply, does a network model that performs well on the training set also perform well on the test set?

Generally speaking, generalization ability comes from learning from diverse data. Therefore:

Stochastic gradient descent: training on one sample at a time usually gives better generalization;
Batch gradient descent: training on the same data every time lacks variety, which limits the improvement of generalization;
Mini-batch gradient descent: a compromise that balances diversity and stability.

Fitting¶

Fitting refers to the performance of the network model on the training set, that is, how well the model has learned the patterns in the data.

Poor generalization is generally caused by an unsatisfactory fit, and there are two typical cases:

Underfitting: the network model performs poorly on the training set, and also performs poorly on the test set. This is usually because of insufficient training, or because the network model is too simple.
Overfitting: the network model performs nearly perfectly on the training set but poorly on the test set. This is usually because the network model is too complex and has learned redundant details that are not representative.

💡 Tip: Mini-batch gradient descent has both good convergence speed and good generalization ability, so it is the most commonly used training method in practice.

import numpy as np

Tensor¶

class Tensor:

    def __init__(self, data):
        self.data = np.array(data)
        self.grad = np.zeros_like(self.data)
        self.gradient_fn = None
        self.parents = set()

    def backward(self):
        if self.gradient_fn is not None:
            self.gradient_fn()

        for p in self.parents:
            p.backward()

    def __str__(self):
        return f'Tensor({self.data})'

Dataset¶

We add an initialization parameter to the dataset: batch_size (the batch size).

By setting different batch sizes, we control how the training is done:

Stochastic gradient descent: this is the dataset’s default (batch_size=1). For our tiny dataset, __getitem__() returns one sample each time, and __len__() returns 4;
Batch gradient descent: simply use all() to get all the data samples at once;
Mini-batch gradient descent: set batch_size to any value greater than 1. For example, with batch_size=2, __getitem__() returns two samples each time, and __len__() returns 2.

Note in particular that for a dataset with batch processing, the output data format changes from the previous 1D tensor with shape (number of features,) to a 2D tensor with shape (batch size, number of features).

class Dataset:

    def __init__(self, batch_size=1):
        self.batch_size = batch_size
        self.load()
        self.train()

    def load(self):
        self.train_data = ([[22.5, 72.0],
                            [31.4, 45.0],
                            [19.8, 85.0],
                            [27.6, 63.0]],
                           [[95],
                            [210],
                            [70],
                            [155]])
        self.test_data = ([[28.1, 58.0]],
                          [[165]])

    def train(self):
        self.data = self.train_data

    def eval(self):
        self.data = self.test_data

    def all(self):
        x, y = self.data
        return Tensor(x), Tensor(y)

    # Number of batches
    def __len__(self):
        x, *_ = self.data
        return len(x) // self.batch_size

    # A single batch
    def __getitem__(self, index):
        s = slice(index * self.batch_size, (index + 1) * self.batch_size)
        x, y = self.data
        return Tensor(x[s]), Tensor(y[s])

Model¶

Once the dataset supports batch processing, the samples returned by __getitem__() are no longer 1D tensors but 2D tensors with shape (batch size, number of features). So the output of the inference function is also a 2D tensor, with shape (batch size, number of labels).

When computing the weight gradient, the gradient function also needs a formula suited to 2D tensors: instead of the multiplication p.grad * x.data, it uses the matrix multiplication p.grad.T @ x.data. When computing the bias gradient, it must also sum along axis 0 (the batch axis): np.sum(p.grad, axis=0).

With these changes, a single computation gives the correct gradient for every sample in the same batch.

class Linear:

    def __init__(self, in_size, out_size):
        self.weight = Tensor(np.ones((out_size, in_size)) / in_size)
        self.bias = Tensor(np.zeros(out_size))

    def __call__(self, x: Tensor):
        return self.forward(x)

    def forward(self, x: Tensor):
        p = Tensor(x.data @ self.weight.data.T + self.bias.data)

        def gradient_fn():
            self.weight.grad += p.grad.T @ x.data
            self.bias.grad += np.sum(p.grad, axis=0)

        p.gradient_fn = gradient_fn
        return p

    @property
    def parameters(self):
        return [self.weight, self.bias]

Loss function (mean squared error)¶

The gradient function needs to divide the error term by the batch size. The purpose is to avoid adding up the error terms of several samples in a batch, which would amplify the gradient several times.

class MSELoss:

    def __call__(self, p: Tensor, y: Tensor):
        return self.loss(p, y)

    def loss(self, p: Tensor, y: Tensor):
        mse = Tensor(np.mean(np.square(y.data - p.data)))

        def gradient_fn():
            p.grad += -2 * (y.data - p.data) / y.data.size

        mse.gradient_fn = gradient_fn
        mse.parents = [p]
        return mse

Optimizer (stochastic gradient descent)¶

class SGDOptimizer:

    def __init__(self, parameters, lr):
        self.parameters = parameters
        self.lr = lr

    def zero_grad(self):
        for p in self.parameters:
            p.grad = np.zeros_like(p.data)

    def step(self):
        for p in self.parameters:
            p.data -= p.grad * self.lr

Hyperparameters¶

Learning rate¶

LEARNING_RATE = 0.00001

Batch size¶

As our second hyperparameter, the batch size defines the number of samples used in each batch.

We will use mini-batch gradient descent. But by adjusting the batch size, we can also get stochastic gradient descent and batch gradient descent.

💡 Tip: In practice, the batch size is usually a power of 2 (such as 16, 32, 64, or 128). This is not a mathematical requirement. Setting it this way simply maximizes hardware efficiency.

BATCH_SIZE = 2

Modeling¶

When creating the dataset, we need to pass the batch size as a parameter.

dataset = Dataset(BATCH_SIZE)
model = Linear(2, 1)
loss_fn = MSELoss()
optimizer = SGDOptimizer(model.parameters, lr=LEARNING_RATE)

Training¶

With mini-batch gradient descent, the 4 samples in the training set need only 2 iterations.

for i in range(len(dataset)):
    feature, label = dataset[i]

    optimizer.zero_grad()
    prediction = model(feature)
    loss = loss_fn(prediction, label)
    loss.backward()
    optimizer.step()

Inference¶

For model inference, we use the full-batch approach.

dataset.eval()
feature, label = dataset.all()

prediction = model(feature)
print(f'prediction:\t{prediction}')

prediction:	Tensor([[56.19176426]])

Evaluation¶

loss = loss_fn(prediction, label)
print(f'loss:\t{loss}')

loss:	Tensor(11839.232164432306)

The evaluation shows that after mini-batch gradient descent reduced the number of iterations and improved efficiency, convergence has slowed down accordingly.

This is understandable. Before, we trained with one sample at a time and updated the parameters 4 times in total. Now, with a batch size of 2, the total number of parameter updates is smaller.

Summary¶

The three approaches compare as follows:

Approach	Batch size	Convergence speed	Generalization	Overfitting risk
Stochastic gradient descent (SGD)	One sample	Slow (large oscillations)	Strong	Low
Batch gradient descent (Batch GD)	All samples	Fast (stable single steps) but demanding on hardware (memory and compute)	Weak	High
Mini-batch gradient descent (Mini-batch GD)	A few samples	Medium	Medium	Medium

Exercises¶

By adjusting the batch size, try batch gradient descent and see how it performs.