Gradient Descent: Batch, Stochastic, and Mini-batch
You can access the code used for this blog post in a Google Colab Notebook here.
Imagine that we are standing on a mountain with a very thick fog which render us with zero visibility. And our goal is to descend the mountain. To descend the mountain, we “feel” the slope of the ground around us and take a step in the direction where the ground slopes down the most.
A more formal definition: Gradient descent is an iterative optimization algorithm used in machine learning and deep learning to minimize a cost function (mountain). It updates model parameters (weights and biases) by calculating the gradient (slope) of the loss function and moving in the opposite direction of that slope, effectively descending toward the lowest error (descending the mountain).
In this blog post, we will look into three variants of the gradient descent algorithm.
Note: For the rest of the blog post and the screenshots, I used the MNIST dataset. Reading my earlier post “Building an Image Classifier from Scratch” can give you more context.
Batch Gradient Descent
In Batch Gradient Descent, we calculate the gradient of the cost function $J(\theta)$ for the entire dataset of size $m$. Let $L$ be the loss function:
$$ \nabla_{\theta}J(\theta) = \frac{1}{m}\sum_{i=1}^{m} \nabla_{\theta}L $$Then we perform the update (descent):
$$ \theta = \theta – \alpha \cdot \nabla_{\theta}J(\theta) $$Where:
- $\theta$ represents the model parameters $W$ (weights) and $b$ (biases).
- $\alpha$ is the learning rate
For example, if we have 100,000 data samples in our training set, it must multiple a weight matrix by the entire dataset matrix BEFORE a single update.
Here’s how the Cost looks like per gradient descent for Batch Gradient Descent.
We can see that learning is rather smooth.
Stochastic Gradient Descent (SGD)
In SGD, we approximate this by picking one random example $(x^{(i)}, y^{(i)})$ at each step. The update rule becomes:
$$ \theta = \theta – \alpha \cdot \nabla_{\theta}L $$Where:
- $\theta$: The model parameters (weights and biases).
- $\alpha$: The learning rate.
- $L$: The loss function for a single observation.
For SGD, we only multiply the weight by ONE training sample at a time.
It is “Stochastic” (random) because a single data point is a “noisy” approximation of the whole dataset.
For SGD, let’s look at two different graphs.
- Cost vs Epoch
2. Cost vs Gradient Descent
The first graph is what the cost look like in every epoch. In the first graph, we can see that the learning process is not smooth. This is even more evident in the second graph, where we look at the cost on every gradient descent. We can see that there’s a lot of “noise”, especially in the early part of training. This is normal because we are updating the gradient in every data sample, so for example, our first training sample is the digit “1”, then the model will correct the weights towards the digit “1”, then if it sees the training sample “0” next, the cost will be higher since visually “1” are all lines without curve so it’s “visually” far from “0”.
Mini-batch SGD
The problem with Batch Gradient Descent is training takes a lot of time because it has to process the entire dataset for every single step.
Mini-batch SGD takes the best of both worlds, instead of being extreme and just taking “one” data example, we get a smaller batch of the dataset. For example, if we have 100,000 data samples in our training set, we take 1024 random samples from it.
The update rule for a single mini-batch $B_k$ is:
$$ \theta = \theta – \alpha \cdot \frac{1}{n} \sum_{i \in B_k} \nabla_{\theta} L $$Where:
- $n$: The batch size (number of samples in the mini-batch).
- $\sum_{i \in B_k}$: The sum of gradients only for the samples in that specific batch.
Again, we also look at two graphs for the Mini-batch:
- Cost vs Epoch
2. Cost vs Gradient Descent
We can see that there are still “noise” but it’s not as much as with SGD.
The Trade-Offs
As in almost all things, choosing which variation to use comes down to trade-off between hardware limitation and convergence speed (amount of time to descend the mountain).
Batch GD requires loading the entire dataset into RAM/VRAM. For MNIST or other smaller datasets, this is fine. But normally, training datasets are hundreds of GBs and loading all of them will require a very powerful and very expensive computer or more often impossible with our current technology.
SGD wastes GPU utilization because we are only computing one sample at a time. Our GPUs are very good at doing matrix computations. Having too much “noise” also often prevents the algorithm to converge, but instead it just “wanders” around.
Mini-batch is the “golden-standard”. The size of the batch is a hyperparameter and can be tuned depending on the hardware being used to train. It also takes advantage of GPUs through “matrix-matrix” multiplications. And the “noise” due to changing of batch often helps our algorithm to get out of being stuck to a local minima.
| Method | Speed per step | Accuracy per step | Memory Usage |
|---|---|---|---|
| Batch | Slow | Perfect slope (Smooth) | High |
| SGD | Fast | Noisy slope (Wanders) | Low |
| Mini-batch | Moderate | Good approximation | Moderate |
Results in MNIST
I also trained an image classifier using the different methods and here are the results.
Batch GD Accuracy: 91.23%
SGD Accuracy: 89.76%
Mini-batch GD Accuracy: 96.58%
For Batch GD, I can’t even use the whole dataset because the training becomes unbearably slow.
The step update in SGD is fast but convergence is slow and as you can see it’s accuracy is the lowest.
Lastly, Mini-batch got the highest accuracy. Mini-batch generalizes better than Batch GD because the slight noise helps it escape bad local minima.
More reads
If you enjoy this post, you might enjoy my other posts:
Full Python Implementation
In case you want to see the code implementation of the three variants of Gradient Descent.
class GradientDescent:
def __init__(self):
np.random.seed(42)
# Layer 1: 784 Inputs -> 128 Hidden Neurons
self.W1 = np.random.randn(128, 784) * 0.01
self.b1 = np.zeros((128, 1))
# Layer 2: 128 Hidden Neurons -> 10 Outputs
self.W2 = np.random.randn(10, 128) * 0.01
self.b2 = np.zeros((10, 1))
self.epsilon = 1e-8
self.cost = []
self.cost_sgd = []
self.cost_mini = []
self.sgd_cost_per_gd = []
self.mini_cost_per_gd = []
def sigmoid(self, z):
# np.clip prevents overflow/underflow errors in exp
z_clipped = np.clip(z, -15, 15)
return 1 / (1 + np.exp(-z_clipped))
def sigmoid_derivative(self, a):
return a * (1 - a)
def softmax(self, z):
max_logit = np.max(z, axis=0, keepdims=True)
capped_z = np.exp(z - max_logit)
return capped_z / np.sum(capped_z, axis=0, keepdims=True)
def forward(self, X):
# Layer 1 (Hidden Layer)
self.Z1 = np.dot(self.W1, X) + self.b1
self.A1 = self.sigmoid(self.Z1)
# Layer 2 (Output Layer)
self.Z2 = np.dot(self.W2, self.A1) + self.b2
self.A2 = self.softmax(self.Z2)
return self.A2
# Batch Gradient Descent
def train(self, X, y, learning_rate=0.01, epochs=1):
m = X.shape[0]
# One-hot encoding
y_one_hot = np.zeros((10, m))
for i in range(m):
y_one_hot[y[i], i] = 1
for epoch in range(epochs):
self.forward(X.T)
# --- Backprop for Layer 2 (Output) ---
dZ2 = self.A2 - y_one_hot
dW2 = (1/m) * np.dot(dZ2, self.A1.T)
db2 = (1/m) * np.sum(dZ2, axis=1, keepdims=True)
# --- Backprop for Layer 1 (Hidden) ---
# We multiply by W2 to see how much the hidden layer contributed to the error
dA1 = np.dot(self.W2.T, dZ2)
dZ1 = dA1 * self.sigmoid_derivative(self.A1)
dW1 = (1/m) * np.dot(dZ1, X)
db1 = (1/m) * np.sum(dZ1, axis=1, keepdims=True)
# Gradient Descent Update
self.W2 = self.W2 - (learning_rate * dW2)
self.b2 = self.b2 - (learning_rate * db2)
self.W1 = self.W1 - (learning_rate * dW1)
self.b1 = self.b1 - (learning_rate * db1)
cost = -np.sum(y_one_hot * np.log(self.A2 + self.epsilon)) / m # adding epsilon prevents log(0)
self.cost.append(cost)
if epoch % 1000 == 0:
print(f"Epoch: {epoch}, Cost: {cost:.4f}")
# Stochastic Gradient Descent
def train_sgd(self, X, y, learning_rate=0.01, epochs=1):
m = X.shape[0]
for epoch in range(epochs):
# 1. Shuffle at the start of every epoch.
indices = np.random.permutation(m)
X_shuffled = X[indices]
y_shuffled = y[indices]
# We want to track the average cost per epoch
epoch_cost = 0
update_ctr = 0
for i in range(m):
# 2. Reshape single image to (784, 1) and label to one-hot (10, 1)
x_i = X_shuffled[i].reshape(-1, 1)
y_i = np.zeros((10, 1))
y_i[y_shuffled[i]] = 1
self.forward(x_i)
# --- Backprop for Layer 2 (Output) ---
dZ2 = self.A2 - y_i
dW2 = np.dot(dZ2, self.A1.T)
db2 = np.sum(dZ2, axis=1, keepdims=True)
# --- Backprop for Layer 1 (Hidden) ---
# We multiply by W2 to see how much the hidden layer contributed to the error
dA1 = np.dot(self.W2.T, dZ2)
dZ1 = dA1 * self.sigmoid_derivative(self.A1)
dW1 = np.dot(dZ1, x_i.T)
db1 = np.sum(dZ1, axis=1, keepdims=True)
# In SGD, we immediately update the weights on each data sample.
self.W2 = self.W2 - (learning_rate * dW2)
self.b2 = self.b2 - (learning_rate * db2)
self.W1 = self.W1 - (learning_rate * dW1)
self.b1 = self.b1 - (learning_rate * db1)
cost = -np.sum(y_i * np.log(self.A2 + self.epsilon)) # adding epsilon prevents log(0)
self.sgd_cost_per_gd.append(cost)
epoch_cost += cost
update_ctr += 1
avg_epoch_cost = epoch_cost / update_ctr
self.cost_sgd.append(avg_epoch_cost)
if epoch % 100 == 0:
print(f"Epoch: {epoch}, Cost: {avg_epoch_cost:.4f}")
# Mini Batch
def train_mini(self, X, y, learning_rate=0.01, epochs=100, batch_size=256):
m = X.shape[0]
for epoch in range(epochs):
# 1. Shuffle at the start of every epoch.
indices = np.random.permutation(m)
X_shuffled = X[indices]
y_shuffled = y[indices]
epoch_cost = 0
update_ctr = 0
for i in range(0, m, batch_size):
# 2. Reshape single image to (784, 1) and label to one-hot (10, 1)
x_batch = X_shuffled[i : i + batch_size]
y_batch_labels = y_shuffled[i : i + batch_size]
n = x_batch.shape[0] # Actual size of the current batch
y_batch = np.zeros((10, n))
for idx, val in enumerate(y_batch_labels):
y_batch[val, idx] = 1
# 3. Forward pass
a2 = self.forward(x_batch.T)
# 4. Backprop
dZ2 = a2 - y_batch
dW2 = (1/n) * np.dot(dZ2, self.A1.T)
db2 = (1/n) * np.sum(dZ2, axis=1, keepdims=True)
# --- Backprop for Layer 1 (Hidden) ---
# We multiply by W2 to see how much the hidden layer contributed to the error
dA1 = np.dot(self.W2.T, dZ2)
dZ1 = dA1 * self.sigmoid_derivative(self.A1)
dW1 = (1/n) * np.dot(dZ1, x_batch)
db1 = (1/n) * np.sum(dZ1, axis=1, keepdims=True)
self.W2 = self.W2 - (learning_rate * dW2)
self.b2 = self.b2 - (learning_rate * db2)
self.W1 = self.W1 - (learning_rate * dW1)
self.b1 = self.b1 - (learning_rate * db1)
# Let's keep track of the cost.
cost = -np.sum(y_batch * np.log(a2 + self.epsilon)) / n
self.mini_cost_per_gd.append(cost)
epoch_cost += cost
update_ctr += 1
avg_epoch_cost = epoch_cost / update_ctr
self.cost_mini.append(avg_epoch_cost)
if epoch % 100 == 0:
print(f"Epoch: {epoch}, Avg Cost: {avg_epoch_cost:.4f}")