One Line, Three Batch Sizes

Linear regression predicts one number from another with a straight line,

\hat{y} = b_0 + b_1 x

and gradient descent finds $b_0$ and $b_1$ by repeatedly stepping downhill on the mean squared error. Batch, stochastic and mini-batch gradient descent are the same algorithm with one difference: how many rows each update looks at. Your job is a single training function that runs all three.

For a batch of $m$ rows, the cost is

J = \frac{1}{m}\sum_{i=1}^{m}\left(\hat{y}_i - y_i\right)^2

and its slopes with respect to the two parameters are

g_0 = \frac{2}{m}\sum_{i=1}^{m}\left(\hat{y}_i - y_i\right), \qquad g_1 = \frac{2}{m}\sum_{i=1}^{m}\left(\hat{y}_i - y_i\right)x_i

The training recipe

Start at $b_0 = 0$ and $b_1 = 0$ .
One epoch is one pass through the data. Walk through the rows in the order given (no shuffling) and cut them into consecutive batches of batch_size rows. If the rows don't divide evenly, the last batch is smaller. If batch_size is bigger than the dataset, the only batch is the whole dataset.
For each batch, compute every prediction $\hat{y}_i$ with the current $b_0$ and $b_1$ . Then compute $g_0$ and $g_1$ , with $m$ equal to the number of rows in that batch. Only then update both parameters, with $\alpha$ = learning_rate:

b_0 \leftarrow b_0 - \alpha\, g_0, \qquad b_1 \leftarrow b_1 - \alpha\, g_1

Task: write fit_line(xs, ys, learning_rate, epochs, batch_size) and return the tuple (b0, b1) after epochs epochs, each rounded to 4 decimal places.

Setting batch_size = len(xs) gives batch gradient descent, batch_size = 1 gives stochastic gradient descent, and anything in between is mini-batch.

Run the same data with all three batch sizes and compare. Stochastic gradient descent makes len(xs) updates per epoch where batch gradient descent makes one, so it travels much further in a single pass. Where it ends up depends on the data. If every row sits exactly on one straight line, all the one-row steps agree and stochastic gradient descent lands on that line. If the rows are scattered, each one-row step tugs the line toward that single row, so with a fixed learning rate it never comes to rest: it repeats the same tugs every epoch, around a point away from the line batch gradient descent closes in on. The bigger the learning rate, the further away that point is.

One Line, Three Batch Sizes

Examples

Hints