SGD and mini-batches

Prerequisites

• Gradient descent.
• Averages of numbers.
• MSE loss.

The analogy

How it works

Per-example loss: $L_i = \tfrac12 (w x_i - y_i)^2$, so its gradient is $g_i = (w x_i - y_i)\,x_i$. A batch gradient is the average of $g_i$ over the batch.

Start with one step, three ways

Points (1, 2), (2, 4), (3, 6); $w = 0$; $\eta = 0.1$. At $w = 0$, each gradient is $g_i = -y_i x_i$. Predict, then open "Check".

Finish the picture: one epoch of SGD

Process all three points once in the order (1, 2), (2, 4), (3, 6), updating after each.

• After (1, 2), then after (2, 4). Predict $w$.
  Check

  After (1, 2): $g = -2$, $w = 0.2$. After (2, 4): $g = (0.4 - 4) \cdot 2 = -7.2$, $w = 0.92$.

• After (3, 6). Predict $w$ and compare with one full-batch step.
  Check

  $g = (2.76 - 6) \cdot 3 = -9.72$, $w = 1.892$. One epoch of SGD made 3 updates and got much closer to 2 than one full-batch step (0.93).

Full worked calculation

One epoch over the three points (1, 2), (2, 4), (3, 6) with the model $\hat y = wx$, loss $\tfrac12(wx - y)^2$, learning rate 0.1, starting at $w = 0$ (the best value is $w = 2$). We compare full-batch, one-point SGD, and mini-batches of 2.

1. 1Gradient for one point

   $$\frac{\partial}{\partial w}\tfrac12(wx - y)^2 = (wx - y)\,x$$

   Prediction error times the input. At $w = 0$ the point (3, 6) gives $(0 - 6) \cdot 3 = -18$, the point (1, 2) only $-2$.

2. 2Full batch: average all three

   $$g = \frac{-2 - 8 - 18}{3} = -9.333, \qquad w = 0 + 0.1 \cdot 9.333 = 0.933$$

   One careful, exact step per epoch. It used all the data, but it only moved $w$ once.

3. 3SGD: first point (2, 4)

   $$g = (0 \cdot 2 - 4) \cdot 2 = -8, \qquad w = 0 + 0.8 = 0.8$$

   One point gives a noisy estimate of the full gradient (−8 instead of −9.33), but we can act on it immediately.

4. 4SGD: second point (1, 2)

   $$g = (0.8 \cdot 1 - 2) \cdot 1 = -1.2, \qquad w = 0.8 + 0.12 = 0.92$$

   This point has a small input, so its gradient is small and the step is tiny. Different points pull with different strength.

5. 5SGD: third point (3, 6)

   $$g = (0.92 \cdot 3 - 6) \cdot 3 = -9.72, \qquad w = 0.92 + 0.972 = 1.892$$

   After one pass SGD is at 1.892, much closer to 2 than full batch (0.933), because it made three updates instead of one.

6. 6Mini-batches of 2: {(1, 2), (3, 6)} then {(2, 4)}

   $$g_1 = \frac{-2 - 18}{2} = -10 \Rightarrow w = 1.0, \qquad g_2 = (1.0 \cdot 2 - 4) \cdot 2 = -4 \Rightarrow w = 1.4$$

   A compromise: less noise than single points, more updates than full batch, and each batch is a matrix product a GPU computes efficiently.

7. 7Compare the losses

   $$L(w) = \tfrac13 \sum \tfrac12 (w x_i - y_i)^2: \quad L(0) = 9.33, \; L(0.933) = 2.66, \; L(1.4) = 0.84, \; L(1.892) = 0.027$$

   Same data, same learning rate, same amount of computation, very different progress. On big datasets this difference is why everybody uses mini-batch SGD.

The rule in plain words

Shuffle the training set each epoch and split it into mini-batches (from 1 example, pure SGD, up to a few thousand).

For each mini-batch: average the gradients of its examples and take one step $w \leftarrow w - \eta \bar g$.

The code

Read this code after you have done the three steps by hand. It will make more sense.

double batch_grad(double x[], double y[], int idx[], int m, double w) {   /* average over a batch */
    double g = 0;
    for (int k = 0; k < m; k++) {
        int i = idx[k];
        g += (w * x[i] - y[i]) * x[i];               /* gradient of 1/2 (wx - y)^2 */
    }
    return g / m;
}

double sgd_step(double x[], double y[], int idx[], int m, double w, double lr) {
    return w - lr * batch_grad(x, y, idx, m, w);
}

flowchart TD
  A["shuffle the training set"] --> B["take the next mini-batch"]
  B --> C["average the gradients of its examples"]
  C --> D["w = w - learning rate × average"]
  D --> E{"more batches this epoch?"}
  E -->|yes| B
  E -->|no| F{"more epochs?"}
  F -->|yes| A
  F -->|no| G["done"]
          

How costly is it?

A mini-batch $B$ estimates the full gradient by averaging; its noise shrinks with the batch size $m$:

Common mistakes

• Confusing epochs with iterations.
• Summing instead of averaging the batch gradients, so the effective learning rate changes with the batch size.
• Increasing the batch size without adjusting the learning rate.

Try it yourself

Did you know?

Key points to remember

If you remember one thing from each line, remember the because.

• SGD is cheap per step because it uses only a few examples to estimate the gradient.
• Mini-batch gradients are noisy because a small sample doesn't represent the whole dataset exactly.
• Larger batches are less noisy because averaging more examples reduces variance.
• Shuffling matters because ordered data biases consecutive updates.
• Many cheap steps beat few exact ones because progress depends more on the number of updates than on their precision.

Self-check questions

Answer these in order, out loud. They follow the page, so your answers become the script of your video.

1. How are full-batch, SGD and mini-batch like polling voters?
   Sample answer

   Asking everyone is exact but slow, one person is fast but noisy, and a few hundred is fast and accurate enough.

2. From $w = 0$ with $\eta = 0.1$, compute the full-batch, SGD (point (2, 4)) and mini-batch ({(1, 2), (3, 6)}) steps.
   Sample answer

   Full batch: $w \approx 0.93$. SGD: $w = 0.8$. Mini-batch: $w = 1.0$.

3. Why does batch_grad divide by m?

   A good answer mentions the learning rate.

   Sample answer

   Averaging keeps the gradient's scale independent of the batch size, so the same learning rate works whatever $m$ is.

4. So what is SGD with mini-batches, in one sentence?
   Sample answer

   Gradient descent where each update uses the average gradient of a small random batch of examples instead of the whole dataset.

5. Why did one epoch of SGD get closer to the answer than one full-batch step?

   A good answer counts updates.

   Sample answer

   SGD made three updates in that epoch, each using the latest $w$, while full-batch made only one.

6. How many updates per epoch, and what does each step cost?
   Sample answer

   $n/m$ updates per epoch, each costing $O(m)$ gradient computations.

7. Why must the data be shuffled each epoch?
   Sample answer

   Otherwise consecutive batches come from similar examples and push the weights in biased directions.

8. If you had to teach this to a friend in 2 minutes, what would you leave out, and why?
   Sample answer

   I'd leave out the $1/\sqrt{m}$ noise formula. I'd keep the polling analogy, the three one-step results from $w = 0$, and "small random batches, many cheap noisy steps, shuffle every epoch".

External resources

• PyTorch docs: torch.optim.SGD: SGD with optional momentum.
• Wikipedia: Stochastic gradient descent: variants and history.
• Google ML Crash Course: Batch size: batch size, epochs and learning rate.

Videos

• YouTube search: StatQuest SGD: StatQuest's clear explanation.