2.6 - Gradient descent: batch, mini-batch and stochastic¶
The notebook Linear regression from scratch trains its model with batch gradient descent: every update uses all the examples. Noah's notes in chapter 2 describe the alternatives and the decisions around them: the batch size (all the examples, one, or a small group), the epochs, the learning rate, the initialization, the stopping criteria, and watching the cost on a validation set as well as on the training set. In this notebook we implement each of them on the same problem, predicting masks from the number of patients:
- Import the libraries.
- The data.
- One training loop for all the variants.
- Batch, mini-batch and stochastic descent.
- The path on the cost surface.
- Why the data is shuffled.
- The learning rate.
- When to stop.
- Where to start.
- Reducing the learning rate when the validation cost stops improving.
Step 1: Import the libraries¶
import numpy as np
import matplotlib.pyplot as plt
np.set_printoptions(precision=3, suppress=True)
Step 2: The data¶
Noah's example uses a training set of 1,000 examples, so we simulate 1,000 days with the same rule as the from-scratch notebook (between 50 and 150 patients, masks $= 25 + 0.8 \times$ patients plus noise) and 250 more days for validation, which are never used to update the parameters. As in that notebook, the number of patients is standardized with the mean and standard deviation of the training days. The least-squares solution gives the lowest cost that any line can reach on these 1,000 days: it is the floor every variant is trying to reach.
rng = np.random.default_rng(42)
def simulate_days(n):
patients = 50 + 100 * rng.random(n)
masks = 25 + 0.8 * patients + rng.normal(0, 10, n)
return patients, masks
x_train, y_train = simulate_days(1000)
x_val, y_val = simulate_days(250)
mean, std = x_train.mean(), x_train.std()
z_train, z_val = (x_train - mean) / std, (x_val - mean) / std # standardized patients
def cost(w, b, z, y):
return np.mean((w * z + b - y) ** 2) # MSE
best_w, best_b = np.polyfit(z_train, y_train, 1)
floor = cost(best_w, best_b, z_train, y_train)
print(f"Training days: {len(y_train)}, validation days: {len(y_val)}")
print(f"Least squares on the scaled data: w = {best_w:.3f}, b = {best_b:.3f}, lowest possible cost = {floor:.3f}")
Training days: 1000, validation days: 250 Least squares on the scaled data: w = 23.905, b = 103.995, lowest possible cost = 102.465
With the patients standardized, $b$ is simply the mean number of masks, 103.995, and $w$ = 23.905 is the change in masks for one standard deviation of patients. No line can do better than a training cost of 102.465 on these days: what is left is the noise we added, whose variance is 100.
Step 3: One training loop for all the variants¶
One epoch is one pass over the training data. In each epoch the loop shuffles the examples, splits them into batches and, for each batch, computes the gradient with the examples of that batch only and updates the parameters. The update is the one in Noah's equations, with $m$ now being the size of the batch:
$$w := w - \alpha \, \frac{2}{m} \sum_{i \in \text{batch}} (\hat{y}_i - y_i)\, z_i \qquad b := b - \alpha \, \frac{2}{m} \sum_{i \in \text{batch}} (\hat{y}_i - y_i)$$
With batch_size=1000 it is batch descent, with batch_size=1 stochastic descent, and anything in between is mini-batch descent. At the end of each epoch the loop records the cost on the training set and, if we pass one, on the validation set. The last two arguments are for step 10.
def gradient(w, b, z, y):
error = w * z + b - y
return 2 * np.mean(error * z), 2 * np.mean(error) # dC/dw, dC/db
def train(z, y, batch_size, learning_rate, epochs, shuffle=True, seed=0, start=(0.0, 0.0),
z_val=None, y_val=None, reduce_factor=None):
order_rng = np.random.default_rng(seed)
w, b = start
m = len(y)
history = {'cost': [], 'val_cost': [], 'learning_rate': [], 'path': [(w, b)], 'updates': 0}
for epoch in range(epochs):
order = order_rng.permutation(m) if shuffle else np.arange(m) # 1. shuffle the data
for first in range(0, m, batch_size): # 2. for each batch
batch = order[first:first + batch_size]
dw, db = gradient(w, b, z[batch], y[batch]) # the gradient of this batch only
w, b = w - learning_rate * dw, b - learning_rate * db
history['path'].append((w, b))
history['updates'] += 1
history['cost'].append(cost(w, b, z, y))
history['learning_rate'].append(learning_rate)
if z_val is not None: # 3. cost on the validation set
history['val_cost'].append(cost(w, b, z_val, y_val))
stalled = len(history['val_cost']) > 1 and history['val_cost'][-1] > history['val_cost'][-2] - 0.01
if reduce_factor is not None and stalled: # 4. step 10: reduce the learning rate
learning_rate = learning_rate * reduce_factor
return w, b, history
Step 4: Batch, mini-batch and stochastic descent¶
We train four times for 20 epochs with the same learning rate, 0.05, changing only the size of the batch. In Noah's example, with 1,000 examples and batches of 100, one epoch takes 10 updates.
runs = {}
for batch_size in [1000, 100, 10, 1]:
w, b, history = train(z_train, y_train, batch_size, learning_rate=0.05, epochs=20)
runs[batch_size] = history
print(f"batch size {batch_size:>4}: {history['updates'] // 20:>4} updates per epoch, "
f"cost after 1 epoch {history['cost'][0]:8.2f}, after 20 epochs {history['cost'][-1]:7.2f}")
plt.figure(figsize=(9, 4))
for batch_size, history in runs.items():
plt.plot(range(1, 21), history['cost'], marker='o', ms=3, label=f'batch size {batch_size}')
plt.axhline(floor, color='gray', ls='--', label='least squares')
plt.yscale('log'); plt.xticks([1, 5, 10, 15, 20]); plt.xlabel('epoch'); plt.ylabel('training cost (MSE)')
plt.legend(); plt.grid(True); plt.show()
batch size 1000: 1 updates per epoch, cost after 1 epoch 9325.54, after 20 epochs 270.77 batch size 100: 10 updates per epoch, cost after 1 epoch 1489.90, after 20 epochs 102.47 batch size 10: 100 updates per epoch, cost after 1 epoch 105.41, after 20 epochs 102.85 batch size 1: 1000 updates per epoch, cost after 1 epoch 138.28, after 20 epochs 130.81
All four runs do the same amount of work per epoch, because each one processes the 1,000 days once; what changes is how many times they update the parameters. Batch descent updates once per epoch, so after 20 epochs it has made only 20 steps and its cost, 270.77, is still far from the floor. Mini-batches of 100 reach it, 102.47, and batches of 10 are already at 105.41 after the first epoch. Stochastic descent is also fast at first, 138.28 after one epoch, but it never settles: each update follows a single day, noise included, so with a constant learning rate its cost jumps up and down from one epoch to the next and ends at 130.81. It is what Noah's notes say: the extremes, all the examples or just one, can be slower or more unstable, and mini-batches give steadier updates without giving up speed.
Step 5: The path on the cost surface¶
The cost is a function of the two parameters, $C(w, b)$, and with scaled data its contour lines are circles around the minimum. On top of them we draw the first updates of three variants: the 20 updates that batch descent makes in 20 epochs, and the first 30 of the other two.
W, B = np.meshgrid(np.linspace(-45, 90, 200), np.linspace(-10, 125, 200))
C = np.mean((W[..., None] * z_train + B[..., None] - y_train) ** 2, axis=-1)
fig, axes = plt.subplots(1, 3, figsize=(16, 5.5), sharey=True)
for ax, batch_size in zip(axes, [1000, 10, 1]):
path = np.array(runs[batch_size]['path'][:31])
ax.contour(W, B, C, levels=[150, 500, 1500, 3500, 6500, 10000], colors='gray', linewidths=0.7)
ax.plot(path[:, 0], path[:, 1], 'o-', ms=3)
ax.plot(best_w, best_b, 'r*', ms=14)
ax.set_aspect('equal')
ax.set_title(f'batch size {batch_size}: first {len(path) - 1} updates')
ax.set_xlabel('w')
axes[0].set_ylabel('b')
plt.tight_layout(); plt.show()
Batch descent walks in a straight line towards the minimum: on circular contours, the direction of steepest descent points straight at the centre. Its steps get shorter as the gradient shrinks, and after 20 updates it has not arrived yet. Mini-batches of 10 follow almost the same direction but wobble, because each batch of 10 days gives a slightly different gradient; in 30 updates, less than a third of an epoch, they are already next to the star. Stochastic descent zigzags, because each step follows one day. It also gets close quickly, but it keeps moving around the minimum instead of stopping on it.
Step 6: Why the data is shuffled¶
Mini-batch and stochastic descent update the parameters before they have seen all the data, so the order in which the examples arrive matters, and real records often come in some order, by date or by size. That is why each epoch of train starts by shuffling them. Suppose the days arrive sorted by the number of patients, from the quietest to the busiest, and we run stochastic descent for one epoch with and without shuffling.
by_patients = np.argsort(x_train)
z_sorted, y_sorted = z_train[by_patients], y_train[by_patients]
for shuffle in [False, True]:
w, b, history = train(z_sorted, y_sorted, batch_size=1, learning_rate=0.05, epochs=1, shuffle=shuffle)
print(f"shuffle={shuffle!s:<5} w = {w:6.3f}, b = {b:7.3f}, cost after one epoch = {history['cost'][-1]:7.2f}")
print(f"least squares: w = {best_w:6.3f}, b = {best_b:7.3f}, cost {floor:.2f}")
shuffle=False w = -24.034, b = 189.562, cost after one epoch = 9722.36 shuffle=True w = 25.701, b = 101.374, cost after one epoch = 112.56 least squares: w = 23.905, b = 103.995, cost 102.46
Without shuffling, one epoch of stochastic descent ends at w = −24.034 and b = 189.562, a line that slopes downwards, with a cost of 9722.36. Each update pulls the line towards the day it has just seen, so at the end of the epoch the line only remembers the busiest days, the last ones it saw, and even the sign of its slope is wrong. With the same days shuffled, the same algorithm ends at w = 25.701 and b = 101.374, with a cost of 112.56, close to the least-squares line. Batch descent does not need shuffling, because each of its updates uses all the days anyway.
Step 7: The learning rate¶
Noah shows the team how training evolves with different learning rates. His figure plots the size of the gradient, which must fall below a tolerance for the algorithm to stop: with a rate that is too large it fluctuates and even grows, because the steps "jump" over the minimum; with one that is too small the line is almost flat. We repeat his experiment with batch descent, 30 updates for each rate, and also follow the bias $b$ on its way to the minimum.
tolerance = 0.01
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 4))
for learning_rate, style in [(0.001, 'o-'), (0.1, 'o-'), (0.9, 'o--'), (1.02, 'o-')]:
w, b, history = train(z_train, y_train, batch_size=1000, learning_rate=learning_rate, epochs=30)
sizes = [np.hypot(*gradient(w, b, z_train, y_train)) for w, b in history['path']]
ax1.plot(sizes, style, ms=3, label=f'learning rate {learning_rate}')
ax2.plot([b for w, b in history['path']], style, ms=3)
print(f"learning rate {learning_rate:<5}: gradient size {sizes[0]:.1f} at the start, {sizes[-1]:10.4f} after 30 updates; "
f"b after 1, 2, 3 updates: {history['path'][1][1]:7.1f} {history['path'][2][1]:7.1f} {history['path'][3][1]:7.1f}")
ax1.axhline(tolerance, color='gray', ls='--', label='tolerance')
ax1.set_yscale('log'); ax1.set_xlabel('update'); ax1.set_ylabel('gradient size |∇C|'); ax1.set_title('The size of the gradient')
ax2.axhline(best_b, color='gray', ls=':', label='minimum of b'); ax2.set_ylim(-150, 250)
ax2.set_xlabel('update'); ax2.set_ylabel('b'); ax2.set_title('The bias on its way to the minimum')
fig.legend(*[sum(h, []) for h in zip(ax1.get_legend_handles_labels(), ax2.get_legend_handles_labels())], loc='lower center', ncol=6)
plt.tight_layout(rect=(0, 0.08, 1, 1)); plt.show()
learning rate 0.001: gradient size 213.4 at the start, 200.9748 after 30 updates; b after 1, 2, 3 updates: 0.2 0.4 0.6 learning rate 0.1 : gradient size 213.4 at the start, 0.2642 after 30 updates; b after 1, 2, 3 updates: 20.8 37.4 50.7 learning rate 0.9 : gradient size 213.4 at the start, 0.2642 after 30 updates; b after 1, 2, 3 updates: 187.2 37.4 157.2 learning rate 1.02 : gradient size 213.4 at the start, 692.1902 after 30 updates; b after 1, 2, 3 updates: 212.2 -8.5 221.0
The three behaviours of Noah's figure appear. With 0.001 the line is almost flat: the gradient goes from 213.4 to 200.9748 in 30 updates and $b$ has barely started to climb, so it would need thousands of updates. With 0.1 the gradient shrinks by the same factor at every update, down to 0.2642, and $b$ climbs smoothly to the minimum. With 1.02 every step lands farther away than the one before, $b$ goes to 212.2, −8.5 and 221.0, and the gradient grows to 692.1902: the algorithm diverges.
The rate of 0.9 shows the jump over the minimum: $b$ goes to 187.2, 37.4 and 157.2, from one side of its minimum, 104.0, to the other. It still converges, because each jump is shorter than the one before. Its gradient has exactly the same size as with 0.1 at every update, and the dashed green line covers the orange one. The reason is that, with this cost and scaled data, each update multiplies the distance to the minimum by $1 - 2\alpha$: 0.8 with a rate of 0.1 and −0.8 with 0.9, the same size but on the other side. For the same reason, any rate above 1 diverges here. The size of the gradient cannot tell 0.1 and 0.9 apart; the parameters can.
Step 8: When to stop¶
Noah lists three stopping criteria: the cost no longer improves significantly, the magnitude of the gradient is very small, or a maximum number of epochs is reached. With batch descent and a learning rate of 0.1 we count how many updates each criterion needs, and how far from the floor the cost is when it fires.
def batch_descent_until(stop, learning_rate=0.1, max_updates=1000):
w, b, previous = 0.0, 0.0, None
for update in range(1, max_updates + 1):
dw, db = gradient(w, b, z_train, y_train)
w, b = w - learning_rate * dw, b - learning_rate * db
current = cost(w, b, z_train, y_train)
if stop(previous, current, np.hypot(dw, db)):
return update, current
previous = current
return max_updates, current
criteria = {
'cost change below 0.001': lambda prev, cur, grad: prev is not None and abs(prev - cur) < 0.001,
'gradient size below 0.01': lambda prev, cur, grad: grad < tolerance,
'maximum of 50 updates': lambda prev, cur, grad: False,
}
for name, stop in criteria.items():
updates, final = batch_descent_until(stop, max_updates=50 if 'maximum' in name else 1000)
print(f"{name:<26} stops after {updates:>3} updates, cost {final:.4f} ({final - floor:+.4f} above the floor)")
cost change below 0.001 stops after 36 updates, cost 102.4659 (+0.0012 above the floor) gradient size below 0.01 stops after 46 updates, cost 102.4647 (+0.0000 above the floor) maximum of 50 updates stops after 50 updates, cost 102.4647 (+0.0000 above the floor)
Stopping when the cost changes by less than 0.001, the rule of the from-scratch notebook, ends after 36 updates, 0.0012 above the floor. Stopping when the gradient is smaller than the tolerance of step 7, 0.01, takes 46 updates and lands on the floor to four decimals. A maximum of 50 updates also gets there, but only because we chose it knowing the answer: with a learning rate of 0.001, 50 updates would stop far from the minimum. That is why the maximum is used as a safety net, together with one of the other two criteria.
Step 9: Where to start¶
Noah recommends random or heuristic initializations, to avoid slow convergence or local minima. We compare three starting points for mini-batch descent (batches of 100): zeros, a random point, and a heuristic one, $b$ equal to the mean number of masks, which is the best constant prediction.
starts = {'zeros': (0.0, 0.0),
'random': tuple(np.random.default_rng(3).normal(0, 50, 2)),
'heuristic (b = mean masks)': (0.0, y_train.mean())}
for name, start in starts.items():
w, b, history = train(z_train, y_train, batch_size=100, learning_rate=0.05, epochs=20, start=start)
first_close = next(e for e, c in enumerate(history['cost'], 1) if c < floor + 0.5)
print(f"{name:<27} start (w, b) = ({start[0]:6.2f}, {start[1]:7.2f}) -> w = {w:.3f}, b = {b:.3f}, "
f"within 0.5 of the floor after {first_close} epoch(s)")
zeros start (w, b) = ( 0.00, 0.00) -> w = 23.849, b = 103.998, within 0.5 of the floor after 5 epoch(s) random start (w, b) = (102.05, -127.78) -> w = 23.849, b = 103.998, within 0.5 of the floor after 6 epoch(s) heuristic (b = mean masks) start (w, b) = ( 0.00, 104.00) -> w = 23.849, b = 103.998, within 0.5 of the floor after 4 epoch(s)
The three starts end on the same line, w = 23.849 and b = 103.998, because this cost has a single minimum: its surface is a bowl, the circles of step 5, and every path leads to the bottom. They even end on the same decimals, because they use the same shuffled batches and after 20 epochs the starting point has been forgotten. The start only changes the length of the way: from the random start, (102.05, −127.78), it takes 6 epochs to come within 0.5 of the floor, from zeros 5, and from the heuristic start 4, because it begins with the bias already right. Noah's warning about local minima is for costs with several valleys, such as those of neural networks; linear regression with the MSE has only one.
Step 10: Reducing the learning rate when the validation cost stops improving¶
Noah's notes insist on monitoring the cost on the validation set as well as on the training set, and mention adaptive methods such as Momentum, RMSProp or Adam, which adjust the learning rate during training. The simplest version of that idea uses the validation cost: after each epoch, if it has not improved by at least 0.01, the learning rate is multiplied by a reduction factor, here 0.5. We compare it with the same stochastic descent at a constant rate. The 250 validation days are never used to update the parameters, only to decide when to reduce the rate.
_, _, constant = train(z_train, y_train, batch_size=1, learning_rate=0.05, epochs=20, z_val=z_val, y_val=y_val)
_, _, reduced = train(z_train, y_train, batch_size=1, learning_rate=0.05, epochs=20, z_val=z_val, y_val=y_val, reduce_factor=0.5)
print(f"constant rate: training cost {constant['cost'][-1]:.3f}, validation cost {constant['val_cost'][-1]:.3f}")
print(f"reduced rate: training cost {reduced['cost'][-1]:.3f}, validation cost {reduced['val_cost'][-1]:.3f}")
print("learning rate in each epoch:", ", ".join(f"{rate:g}" for rate in reduced['learning_rate']))
best_val = cost(best_w, best_b, z_val, y_val)
print(f"validation cost of the least-squares line: {best_val:.3f}")
fig, ax = plt.subplots(figsize=(9, 4))
ax.plot(range(1, 21), constant['val_cost'], marker='o', ms=3, label='constant learning rate')
ax.plot(range(1, 21), reduced['val_cost'], marker='o', ms=3, label='halved when validation stops improving')
ax.axhline(best_val, color='gray', ls='--', label='least-squares line')
ax.set_xticks([1, 5, 10, 15, 20]); ax.set_xlabel('epoch'); ax.set_ylabel('validation cost'); ax.legend(); ax.grid(True); plt.show()
constant rate: training cost 130.810, validation cost 136.984 reduced rate: training cost 102.468, validation cost 122.273 learning rate in each epoch: 0.05, 0.05, 0.05, 0.025, 0.0125, 0.0125, 0.00625, 0.003125, 0.003125, 0.003125, 0.003125, 0.0015625, 0.0015625, 0.0015625, 0.00078125, 0.000390625, 0.000195313, 0.000195313, 9.76563e-05, 9.76563e-05 validation cost of the least-squares line: 122.092
With a constant rate, stochastic descent ends the 20 epochs with a training cost of 130.810 and a validation cost of 136.984, and its validation cost jumps up and down in the figure: each update follows a single day, and at 0.05 the line never settles. Halving the rate whenever the validation cost stalls makes the steps smaller and smaller, from 0.05 to 9.76563e-05, and the line settles: a training cost of 102.468, practically the floor, and a validation cost of 122.273, next to the 122.092 of the least-squares line. The validation days have their own noise, so even the best line scores higher on them than on the training days.
Conclusion¶
- The same update rule gives three algorithms depending on how many examples each step uses. Mini-batches are the usual compromise: many updates per epoch, each one cheap and not too noisy.
- One epoch is always one pass over the data, but the number of updates in it is $m$ divided by the batch size: 1, 10, 100 or 1,000 in this notebook.
- Shuffling matters for any variant that updates before seeing all the data.
- The learning rate, the stopping criterion and the starting point are decisions, and each one can be checked with a few lines of code.
- Watching the cost on a validation set, and reducing the learning rate when it stops improving, is how stochastic methods settle near the minimum instead of bouncing around it.