A network is almost never shown one example at a time, and almost never shown all of them at once. It is shown batches. The batch size decides how noisy each gradient is, how much memory the run needs, and - through normalisation layers - sometimes what the model predicts. It is the hyperparameter people tune last and should tune early.
The vocabulary, precisely
- A batch is the group of examples whose gradients are averaged into one weight update.
- A step (or iteration) is one weight update: one batch forward, one backward, one update.
- An epoch is one pass over the whole training set. With 50,000 examples and a batch of 100, an epoch is 500 steps.
Those definitions matter because the two numbers people compare runs by - epochs and steps - are not interchangeable the moment the batch size changes.
import pandas as pd
examples = 50000
rows = []
for batch in [1, 32, 256, 2048, 50000]:
steps_per_epoch = -(-examples // batch) # ceiling division
rows.append({'batch size': batch,
'steps per epoch': steps_per_epoch,
'updates in 10 epochs': steps_per_epoch * 10,
'name': ('stochastic' if batch == 1 else
'full batch' if batch >= examples else
'mini-batch')})
print(pd.DataFrame(rows).to_string(index=False))
print()
print('Ten epochs at a batch of 32 is 15,630 updates. Ten epochs')
print('at a batch of 2048 is 250. Reporting "trained for 10 epochs"')
print('without the batch size describes two runs that differ by a')
print('factor of sixty in how much learning happened.')
Open this example in the editor and change it. Nothing you type is saved.
Why batches at all
Two reasons, pulling in the same direction.
The first is memory. The activations of every layer for every example in the batch have to be held until the backward pass uses them, because the gradient of a weight needs the activation that it multiplied.
import pandas as pd
# A small convolutional network, activations only - no weights.
layers = [('input 3x224x224', 3 * 224 * 224),
('conv1 64x112x112', 64 * 112 * 112),
('conv2 128x56x56', 128 * 56 * 56),
('conv3 256x28x28', 256 * 28 * 28),
('conv4 512x14x14', 512 * 14 * 14)]
per_example = sum(size for _, size in layers)
rows = []
for batch in [1, 8, 32, 128, 512]:
bytes_needed = per_example * batch * 4 # float32
rows.append({'batch size': batch,
'activation memory': f'{bytes_needed / 1024 ** 2:,.0f} MiB'})
print(f'activations per example: {per_example:,} floats, '
f'{per_example * 4 / 1024 ** 2:.1f} MiB')
print(pd.DataFrame(rows).to_string(index=False))
print()
print('Memory is linear in the batch size, and it is the')
print('activations that dominate, not the weights. This is the')
print('whole reason "CUDA out of memory" is answered by halving')
print('the batch size and nothing else.')
Open this example in the editor and change it. Nothing you type is saved.
The second reason is that an average of many gradients is a better estimate of the true gradient than one example's. The improvement follows the square root of the batch size, which is the fact that governs every batch-size decision.
import numpy as np
import pandas as pd
rng = np.random.default_rng(0)
# The true gradient of some parameter is 1.0; each example gives a
# noisy reading of it.
TRUE = 1.0
NOISE = 3.0
rows = []
for batch in [1, 4, 16, 64, 256, 1024]:
estimates = rng.normal(TRUE, NOISE, (4000, batch)).mean(axis=1)
rows.append({'batch size': batch,
'std of the estimate': round(float(estimates.std()), 4),
'predicted (noise/sqrt(n))': round(NOISE / np.sqrt(batch), 4),
'wrong sign': f'{float((estimates < 0).mean()):.1%}'})
print(pd.DataFrame(rows).to_string(index=False))
print()
print('Measured and predicted agree to three decimals. Note the')
print('last column: at a batch of 1, more than a third of the')
print('updates point the wrong way. At 64 it is under 1%.')
print()
print('And note the cost. Going from 16 to 1024 is 64 times the')
print('compute for 8 times less noise. Beyond a few hundred, you')
print('are buying very little accuracy very expensively, which is')
print('why batch sizes cluster where they do.')
Open this example in the editor and change it. Nothing you type is saved.
The large-batch bargain, and its catch
A larger batch gives a less noisy gradient, so it tolerates a larger learning rate. The rule of thumb is linear scaling: double the batch, double the rate. It holds over a useful range and then stops, and the stopping point is worth seeing rather than reading about.
import numpy as np
import pandas as pd
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
rng = np.random.default_rng(0)
X, y = make_classification(n_samples=12000, n_features=20, n_informative=12,
n_redundant=4, class_sep=0.8, random_state=0)
y = y.reshape(-1, 1).astype(float)
X = (X - X.mean(axis=0)) / X.std(axis=0)
X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.25, random_state=0)
def sigmoid(z):
return 1 / (1 + np.exp(-np.clip(z, -60, 60)))
def train(batch_size, learning_rate, epochs=12, seed=0, width=48):
generator = np.random.default_rng(seed)
W1 = generator.normal(0, np.sqrt(2 / X_tr.shape[1]), (X_tr.shape[1], width))
b1 = np.zeros(width)
W2 = generator.normal(0, np.sqrt(2 / width), (width, 1))
b2 = np.zeros(1)
steps = 0
for _ in range(epochs):
order = generator.permutation(len(X_tr))
for start in range(0, len(X_tr), batch_size):
index = order[start:start + batch_size]
xb, yb = X_tr[index], y_tr[index]
z1 = xb @ W1 + b1
a1 = np.maximum(0, z1)
a2 = sigmoid(a1 @ W2 + b2)
delta2 = (a2 - yb) / len(xb)
gW2, gb2 = a1.T @ delta2, delta2.sum(axis=0)
delta1 = (delta2 @ W2.T) * (z1 > 0)
gW1, gb1 = xb.T @ delta1, delta1.sum(axis=0)
W1 -= learning_rate * gW1
b1 -= learning_rate * gb1
W2 -= learning_rate * gW2
b2 -= learning_rate * gb2
steps += 1
predicted = sigmoid(np.maximum(0, X_te @ W1 + b1) @ W2 + b2)
return float(((predicted > 0.5) == y_te).mean()), steps
print('twelve epochs each; the rate scaled linearly with the batch:')
rows = []
for batch, rate in [(32, 0.05), (64, 0.10), (128, 0.20),
(512, 0.80), (2048, 3.20), (8192, 12.80)]:
accuracy, steps = train(batch, rate)
rows.append({'batch': batch, 'rate': rate, 'updates': steps,
'test accuracy': round(accuracy, 4)})
print(pd.DataFrame(rows).to_string(index=False))
print()
print('Linear scaling holds all the way to a batch of 2048: the')
print('accuracy is flat - in fact it creeps up - while the number')
print('of updates falls from 3,384 to 60. That is the bargain, and')
print('it is a real one. At 8192 it breaks completely: 24 updates')
print('at a rate of 12.8 collapses to 48.5%, which is a coin toss.')
print('The range over which it holds is a property of the problem,')
print('not a universal constant, so the sweep is the only way to')
print('find where your own breaks.')
Open this example in the editor and change it. Nothing you type is saved.
So the catch is not subtle. A large batch means fewer updates per epoch, and at some point there are simply not enough updates left, however large each one is. The industrial answer is to train for more epochs at the large batch - which spends the compute the large batch was supposed to save.
Normalisation
Normalisation is the other half of this lesson, and it is tied to the batch size in a way that surprises people.
Start with the input. Features on wildly different scales make the loss surface a narrow valley, and the earlier lesson showed what gradient descent does in one of those.
import numpy as np
import pandas as pd
rng = np.random.default_rng(1)
n = 5000
age = rng.uniform(18, 80, n) # tens
income = rng.uniform(15000, 200000, n) # tens of thousands
clicks = rng.uniform(0, 1, n) # fractions
raw = np.column_stack([age, income, clicks])
standardised = (raw - raw.mean(axis=0)) / raw.std(axis=0)
print('column raw mean raw std standardised mean std')
for index, name in enumerate(['age', 'income', 'clicks']):
print(f'{name:<10} {raw[:, index].mean():11.2f} {raw[:, index].std():11.2f}'
f' {standardised[:, index].mean():19.6f} {standardised[:, index].std():5.2f}')
print()
# The curvature of a linear model's loss surface is set by the
# covariance of the inputs. Its condition number is what gradient
# descent has to live with.
for label, matrix in [('raw', raw), ('standardised', standardised)]:
eigenvalues = np.linalg.eigvalsh(np.cov(matrix.T))
print(f'{label:<14} condition number '
f'{eigenvalues.max() / eigenvalues.min():,.0f}')
print()
print('Ten orders of magnitude, reduced to 1 - a perfect sphere,')
print('because these three columns are independent. The raw')
print('version is not slow to train, it is untrainable: the rate')
print('that the income direction needs would diverge in the clicks')
print('direction on the first step.')
Open this example in the editor and change it. Nothing you type is saved.
Then normalise inside the network too, because the same problem recurs at every layer as the weights move.
import numpy as np
rng = np.random.default_rng(2)
batch, features = 8, 4
activations = rng.normal(5.0, 3.0, (batch, features)) * np.array([1, 10, 0.1, 4])
def batch_norm(x, epsilon=1e-5):
# Statistics ACROSS THE BATCH, one pair per feature.
return (x - x.mean(axis=0)) / np.sqrt(x.var(axis=0) + epsilon)
def layer_norm(x, epsilon=1e-5):
# Statistics ACROSS THE FEATURES, one pair per example.
return ((x - x.mean(axis=1, keepdims=True))
/ np.sqrt(x.var(axis=1, keepdims=True) + epsilon))
print('input, per-feature std:', activations.std(axis=0).round(2))
print()
print('after batch norm, per-feature std :', batch_norm(activations).std(axis=0).round(4))
print('after batch norm, per-example std :', batch_norm(activations).std(axis=1).round(4))
print()
print('after layer norm, per-feature std :', layer_norm(activations).std(axis=0).round(4))
print('after layer norm, per-example std :', layer_norm(activations).std(axis=1).round(4))
print()
print('Batch norm makes every FEATURE unit variance; layer norm')
print('makes every EXAMPLE unit variance. They normalise across')
print('different axes of the same matrix, and that single')
print('difference is why one depends on the batch and the other')
print('does not.')
Open this example in the editor and change it. Nothing you type is saved.
| Batch normalisation | Layer normalisation | |
|---|---|---|
| Statistics over | the batch, per feature | the features, per example |
| Depends on batch size | yes, badly below about 16 | no |
| Train and eval differ | yes - running averages at eval | no |
| Each example's output depends on the others in its batch | yes | no |
| Standard for | convolutional networks | transformers, recurrent networks |
The fourth row is the one that causes production incidents. Under batch normalisation, during training, what a network predicts for one image depends on which other images happened to share its batch.
import numpy as np
rng = np.random.default_rng(3)
# One example, put through batch norm beside two different sets of
# companions. Nothing about the example changed.
example = np.array([[2.0, -1.0, 4.0]])
def batch_norm(x, epsilon=1e-5):
return (x - x.mean(axis=0)) / np.sqrt(x.var(axis=0) + epsilon)
quiet_companions = rng.normal(0, 0.5, (7, 3))
loud_companions = rng.normal(0, 8.0, (7, 3))
with_quiet = batch_norm(np.vstack([example, quiet_companions]))[0]
with_loud = batch_norm(np.vstack([example, loud_companions]))[0]
print('the example alongside quiet companions:', with_quiet.round(3))
print('the example alongside loud companions :', with_loud.round(3))
print()
print('Same example, same weights, two different activations. At')
print('training time this is a feature - it acts as a regulariser.')
print('At inference time it would be a bug, which is why batch')
print('norm keeps running averages and uses those in eval mode.')
print('Forgetting model.eval() is the most common serving bug in')
print('all of deep learning, and this is the mechanism.')
Open this example in the editor and change it. Nothing you type is saved.
And here is the batch-size interaction, measured. Batch norm estimates its statistics from the batch, so a small batch estimates them badly.
import numpy as np
import pandas as pd
rng = np.random.default_rng(4)
TRUE_MEAN, TRUE_STD = 3.0, 2.0
population = rng.normal(TRUE_MEAN, TRUE_STD, 200000)
rows = []
for batch in [2, 4, 8, 16, 32, 128]:
samples = rng.choice(population, (3000, batch))
estimated_means = samples.mean(axis=1)
estimated_stds = samples.std(axis=1)
rows.append({'batch size': batch,
'mean estimate, std of': round(float(estimated_means.std()), 3),
'std estimate, average': round(float(estimated_stds.mean()), 3),
'bias in the std': round(float(estimated_stds.mean() - TRUE_STD), 3)})
print(f'true mean {TRUE_MEAN}, true std {TRUE_STD}')
print(pd.DataFrame(rows).to_string(index=False))
print()
print('At a batch of 2 the standard deviation is underestimated by')
print('0.8 - a 40% error - and the mean estimate wobbles by 1.4.')
print('Batch norm then divides by that wrong number, every step.')
print('This is why a model that trains fine at batch 64 falls apart')
print('at batch 2 on a smaller GPU, and why group normalisation')
print('exists.')
Open this example in the editor and change it. Nothing you type is saved.
Gradient accumulation: a large batch on a small card
The fix for "I need a batch of 512 and 64 is all that fits" is to run eight batches of 64, add up their gradients, and update once. It is arithmetically identical to a batch of 512 - with one exception, which is the point of the example.
import numpy as np
rng = np.random.default_rng(5)
features = 6
weights = rng.normal(0, 1, features)
X = rng.normal(0, 1, (512, features))
y = X @ np.array([1.0, -2.0, 0.5, 0.0, 3.0, -1.0]) + rng.normal(0, 0.1, 512)
def gradient(X, y, weights):
residual = X @ weights - y
return 2 * X.T @ residual / len(X)
one_big_batch = gradient(X, y, weights)
# Eight micro-batches of 64, accumulated. The MEAN of the eight
# gradients, not the sum - the loss was already a mean over 64.
accumulated = np.zeros(features)
for start in range(0, 512, 64):
accumulated += gradient(X[start:start + 64], y[start:start + 64], weights)
accumulated /= 8
print('batch of 512 ', one_big_batch.round(8))
print('8 x 64, averaged ', accumulated.round(8))
print('identical:', bool(np.allclose(one_big_batch, accumulated)))
print()
# The mistake: summing instead of averaging multiplies the
# effective learning rate by the number of micro-batches.
summed = accumulated * 8
print('8 x 64, summed ', summed.round(8))
print(f'that is {np.abs(summed).max() / np.abs(one_big_batch).max():.0f}x '
f'the correct gradient, so an 8x learning rate you did not ask for')
print()
print('What accumulation does NOT reproduce is batch')
print('normalisation: its statistics are still computed over 64')
print('examples, because that is what is in memory at once. A run')
print('using batch norm and accumulation has the batch size of the')
print('micro-batch for normalisation and of the accumulated batch')
print('for the update, and the two are no longer the same number.')
Open this example in the editor and change it. Nothing you type is saved.
A worked example: one dataset, the whole grid
import numpy as np
import pandas as pd
from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split
X, y = make_moons(n_samples=6000, noise=0.22, random_state=0)
y = y.reshape(-1, 1).astype(float)
X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.25, random_state=0)
X_mean, X_std = X_tr.mean(axis=0), X_tr.std(axis=0)
X_tr_n, X_te_n = (X_tr - X_mean) / X_std, (X_te - X_mean) / X_std
def sigmoid(z):
return 1 / (1 + np.exp(-np.clip(z, -60, 60)))
def loss_of(predicted, actual):
p = np.clip(predicted, 1e-9, 1 - 1e-9)
return float(-np.mean(actual * np.log(p) + (1 - actual) * np.log(1 - p)))
# One column in metres and one in micrometres - the ordinary case
# of two features in different units.
SCALE = np.array([1.0, 500.0])
X_tr_s, X_te_s = X_tr * SCALE, X_te * SCALE
S_mean, S_std = X_tr_s.mean(axis=0), X_tr_s.std(axis=0)
X_tr_sn, X_te_sn = (X_tr_s - S_mean) / S_std, (X_te_s - S_mean) / S_std
def run(batch_size, learning_rate, epochs, inputs='normalised',
width=32, seed=0):
train_X, test_X = {
'normalised': (X_tr_n, X_te_n),
'badly scaled': (X_tr_s, X_te_s),
'badly scaled, then standardised': (X_tr_sn, X_te_sn),
}[inputs]
generator = np.random.default_rng(seed)
W1 = generator.normal(0, np.sqrt(2 / 2), (2, width))
b1 = np.zeros(width)
W2 = generator.normal(0, np.sqrt(2 / width), (width, 1))
b2 = np.zeros(1)
steps = 0
for _ in range(epochs):
order = generator.permutation(len(train_X))
for start in range(0, len(train_X), batch_size):
index = order[start:start + batch_size]
xb, yb = train_X[index], y_tr[index]
z1 = xb @ W1 + b1
a1 = np.maximum(0, z1)
a2 = sigmoid(a1 @ W2 + b2)
delta2 = (a2 - yb) / len(xb)
gW2, gb2 = a1.T @ delta2, delta2.sum(axis=0)
delta1 = (delta2 @ W2.T) * (z1 > 0)
gW1, gb1 = xb.T @ delta1, delta1.sum(axis=0)
W1 -= learning_rate * gW1
b1 -= learning_rate * gb1
W2 -= learning_rate * gW2
b2 -= learning_rate * gb2
steps += 1
if not np.isfinite(W1).all():
return None, steps
predicted = sigmoid(np.maximum(0, test_X @ W1 + b1) @ W2 + b2)
return float(((predicted > 0.5) == y_te).mean()), steps
# ---- 1. Batch size at a FIXED rate: noise, and nothing else.
print('batch size at a fixed rate of 0.1, 20 epochs:')
rows = []
for batch in [1, 8, 32, 128, 512, 4500]:
accuracy, steps = run(batch, 0.1, 20)
rows.append({'batch': batch, 'updates': steps,
'test accuracy': 'diverged' if accuracy is None else round(accuracy, 4)})
print(pd.DataFrame(rows).to_string(index=False))
print('-> the rate was tuned for the small batches, so the large')
print(' ones are starved of updates. Nothing has been held')
print(' constant except a number that should have moved.')
print()
# ---- 2. The same batch sizes with the rate scaled.
print('the same batch sizes, rate scaled linearly from 0.025 at batch 8:')
rows = []
for batch in [8, 32, 128, 512, 4500]:
rate = 0.025 * batch / 8
accuracy, steps = run(batch, rate, 20)
rows.append({'batch': batch, 'rate': round(rate, 4), 'updates': steps,
'test accuracy': 'diverged' if accuracy is None else round(accuracy, 4)})
print(pd.DataFrame(rows).to_string(index=False))
print('-> the same five batch sizes, now within half a point of')
print(' each other up to 512, on a fiftieth of the updates.')
print()
# ---- 3. Input normalisation, on and off, on the same grid.
print('input scaling, at batch 32 and rate 0.1:')
rows = []
for inputs in ['normalised', 'badly scaled', 'badly scaled, then standardised']:
accuracy, _ = run(32, 0.1, 20, inputs=inputs)
rows.append({'inputs': inputs,
'test accuracy': 'diverged' if accuracy is None else round(accuracy, 4)})
print(pd.DataFrame(rows).to_string(index=False))
print('-> multiplying one column by 500 is enough to stop the')
print(' network learning at all. Standardising it afterwards')
print(' recovers the original result exactly, because')
print(' standardisation undoes the scaling.')
print()
# ---- 4. Epochs at two batch sizes, to show what "an epoch" hides.
print('accuracy against epochs, both at rate 0.1:')
rows = []
for epochs in [1, 3, 10, 30]:
small, small_steps = run(32, 0.1, epochs)
large, large_steps = run(512, 0.1, epochs)
rows.append({'epochs': epochs,
'batch 32': round(small, 4),
'updates': small_steps,
'batch 512': round(large, 4),
'updates ': large_steps})
print(pd.DataFrame(rows).to_string(index=False))
print('-> at thirty epochs the large batch has had 270 updates and')
print(' sits at 0.896, which the small batch passed inside its')
print(' first three. "Thirty epochs" described both runs.')
Open this example in the editor and change it. Nothing you type is saved.
Read the four tables together and the lesson is one sentence: the batch size is not a performance knob, it is a coupled change to the learning rate and the number of updates. Change it on its own and the run gets worse in one of two ways - too noisy, or too few steps. Change it with the rate and it is nearly free, over a range.
Failure modes
| Symptom | Cause |
|---|---|
CUDA out of memory partway through the first epoch |
A batch with longer sequences or larger images than the earlier ones; sort by length or cap it |
| Loss jagged with no trend | Batch too small for the learning rate |
| Loss falls smoothly but the final result is worse than a smaller batch | Too few updates; train more epochs or lower the batch |
| Works at batch 64, diverges at batch 8 | The rate was tuned for 64; scale it down with the batch |
| Validation much worse than training from epoch one, model uses batch norm | model.eval() not called, so batch statistics are being used at evaluation |
| Trains fine, serves badly on single requests | Batch norm with a batch of one at inference; eval mode, or layer norm |
| Accumulation "works" but the loss explodes | Gradients summed instead of averaged, so an 8x rate |
| Results change between runs with everything fixed | The shuffle seed; batch composition is part of the result |
Check your understanding
The gradient of a mean is the mean of the gradients, so accumulation must be exact. Verify it rather than believing it - and verify the one case where the equivalence genuinely breaks:
import numpy as np
rng = np.random.default_rng(0)
X = rng.normal(0, 1, (240, 5))
y = rng.normal(0, 1, 240)
weights = rng.normal(0, 1, 5)
def gradient(X, y, w):
return 2 * X.T @ (X @ w - y) / len(X)
full = gradient(X, y, weights)
accumulated = sum(gradient(X[i:i + 60], y[i:i + 60], weights)
for i in range(0, 240, 60)) / 4
print('gradients agree:', bool(np.allclose(full, accumulated)))
# Now the case that does not survive: statistics computed over the
# batch rather than per example.
def batch_norm_mean(X):
return X.mean(axis=0)
print('full-batch means ', batch_norm_mean(X).round(4))
print('micro-batch means ')
for i in range(0, 240, 60):
print(' ', batch_norm_mean(X[i:i + 60]).round(4))
print()
print('The gradient accumulates exactly. The normalisation')
print('statistics do not, and no amount of accumulation recovers')
print('them. If your model normalises over the batch, the')
print('micro-batch size is a real hyperparameter, not an')
print('implementation detail.')
Open this example in the editor and change it. Nothing you type is saved.