Same Data, Different Handfuls

Batch size is a dial between two costs. Small batches give many updates per pass through the data, but each one points in a wobbly direction. Large batches give steady directions, but only a few updates. This problem measures both sides of the dial on real numbers.

A one-weight model predicts $\hat{y} = w\,x$ and scores each example by its squared error $(w\,x_i - y_i)^2$. That example's gradient with respect to $w$ is

$g_i = 2\,x_i\,(w\,x_i - y_i)$

To measure the noise without training getting in the way, the weight is frozen at a single value $w$ and never updated. The examples are read in the order given, which is already shuffled.

For one batch size $B$:

• Split the examples into consecutive batches of $B$, in order. The last batch takes whatever is left over, so it can be smaller than $B$. It is never dropped. A batch size at least as large as the dataset gives a single batch that holds everything.
• A batch's gradient $g^{(k)}$ is the average of $g_i$ over the examples actually in that batch. Each batch would produce one update, so the number of batches $K$ is the number of updates in one pass.
• The full-batch gradient $G$ is the average of $g_i$ over all $n$ examples.
• The noise is the root-mean-square gap between the batch gradients and $G$:

$\text{noise} = \sqrt{\frac{1}{K}\sum_{k=1}^{K}\left(g^{(k)} - G\right)^2}$

Task: write batch_size_sweep(xs, ys, w, batch_sizes). For each batch size in batch_sizes, in the order given, return the pair [updates, noise]: the number of batches $K$ as a whole number, and the noise rounded to 4 decimal places. The result is a list of these pairs.

xs and ys have the same length, at least 1. Every batch size is a whole number of at least 1.

Worked through, batch_size_sweep([1.0, 2.0, 3.0], [1.0, 3.0, 2.0], 1.0, [2]):

$G = (0 - 4 + 6)/3 = 0.6667$. With $B = 2$ the batches are examples 1–2, with gradient $(0 - 4)/2 = -2$, and example 3 alone, with gradient $6/1 = 6$. So $K = 2$ and

$\text{noise} = \sqrt{\tfrac{1}{2}\left[(-2 - 0.6667)^2 + (6 - 0.6667)^2\right]} = \sqrt{\tfrac{1}{2}(7.1111 + 28.4444)} = 4.2164$

The answer is [[2, 4.2164]].

Running several batch sizes over the same data puts both costs side by side: how many updates a pass buys, and how far each of those updates strays from the direction the whole dataset would have given.