ML Foundations
Why averaging batch means misweights the last batch
A smaller final batch can distort an epoch loss when every batch mean gets equal weight. Use the denominator that matches the loss you mean to report.
Evidence and testing
Review record: Tensor Almanac automated evidence gate (no independent human review). Example versions, recorded outputs and limitations appear below.
A short batch is still a batch
Training and evaluation loops often collect one scalar loss per batch and print an epoch summary at the end. A common pattern adds each batch mean to a running total, then divides by the number of batches. That calculation gives every batch equal influence, even when the final batch contains fewer examples.
PyTorch's DataLoader keeps the final incomplete batch by default when the dataset size is not divisible by the batch size; setting `drop_last=True` removes it instead. So a four-example dataset with batch size three can produce batches of sizes three and one. If the one-example batch has a very different loss, an unweighted average of the two batch means no longer equals the mean across examples. [2]
References: [2]
Reproduce the weighting difference with four losses
Use four per-example squared-error values, 1, 3, 5 and 100, and split them into batches of three and one. The batch means are 3 and 100. Averaging those two summaries gives 51.5; averaging the four example losses gives 27.25. The arithmetic is small enough to calculate directly and does not require a framework installation.
The script compares an unweighted mean of batch means with a count-weighted result. This is an executable arithmetic model of one scalar loss per example. It is not a PyTorch training run, a performance benchmark, or a change to gradient computation. [1, 2]
Executed with Python 3.12.14 standard library only · 2026-10-22
"""Show why per-batch means need denominator-aware epoch aggregation."""
per_example_losses = [1.0, 3.0, 5.0, 100.0]
batch_size = 3
batches = [
per_example_losses[start:start + batch_size]
for start in range(0, len(per_example_losses), batch_size)
]
batch_means = [sum(batch) / len(batch) for batch in batches]
unweighted_epoch_mean = sum(batch_means) / len(batch_means)
total_loss = sum(sum(batch) for batch in batches)
total_examples = sum(len(batch) for batch in batches)
sample_weighted_epoch_mean = total_loss / total_examples
print("batch sizes:", [len(batch) for batch in batches])
print("batch means:", batch_means)
print(f"unweighted mean of batch means: {unweighted_epoch_mean:.2f}")
print(f"sample-weighted epoch mean: {sample_weighted_epoch_mean:.2f}")
assert batch_means == [3.0, 100.0]
assert unweighted_epoch_mean == 51.5
assert sample_weighted_epoch_mean == 27.25
assert sample_weighted_epoch_mean == sum(per_example_losses) / len(per_example_losses)
Recorded output
batch sizes: [3, 1] batch means: [3.0, 100.0] unweighted mean of batch means: 51.50 sample-weighted epoch mean: 27.25
Weight by the denominator your loss actually uses
For an unweighted scalar loss per example, accumulate each batch mean multiplied by that batch's example count, then divide by the total number of examples. Equivalently, accumulate the per-example losses themselves. In the script, multiplying the first batch mean by three and the second by one reconstructs the total of 109 before division by four.
A library's reduction semantics determine the right denominator. PyTorch documents MSELoss with `reduction='mean'` as dividing the sum by the number of loss elements, and its inputs may have arbitrary shapes. If each example contributes multiple elements, sample count alone only matches the element-wise mean when examples contribute the same number of elements. Masks, sample weights, ignored targets and custom reductions need their own denominator. [1]
References: [1]
Keep reporting separate from optimization
This correction concerns an epoch summary assembled from batch results. It does not mean every training step should be reweighted after computing its gradient: changing loss scaling during optimizer updates can change the optimization procedure. Keep the training objective and update behavior intact unless you intentionally want to change them; use denominator-aware accumulation for the dataset-level statistic you intend to report.
The example assumes four included examples, one scalar unweighted loss per example, and an ordinary arithmetic mean. With `drop_last=True`, the final incomplete batch is discarded and the reported mean covers only the examples that remain; that changes sample coverage. Distributed samplers, duplicated examples, weighted objectives and task-specific aggregation require a statistic defined for that setup. [1, 2]