Gradient descent, overfitting, and regularization in LMs · CS336 Language Modeling From Scratch

Stochastic gradient descent and mini-batching

You rarely can (or should) compute gradients on your entire dataset every step; instead you use stochastic gradient descent (SGD) with mini-batches, which trades exactness for speed.

Full-batch vs mini-batch gradient descent

Full-batch gradient descent:

At each step, you use all $N$ training examples to compute the gradient.
Gradient is the exact derivative of loss over the training set.
Updates are stable and low-variance but very expensive per step.

Mini-batch SGD:

At each step, you sample a batch of $B \ll N$ examples.
Gradient is an estimate of the full gradient based on this batch.
Updates are cheap per step and much more frequent, but noisy.

Mechanically, for parameters $\theta$ and loss $L(\theta)$ :

Full batch (size $N$ ):
$g_{\text{full}} = \frac{1}{N}\sum_{i=1}^{N} \nabla_\theta \ell(x_i, y_i; \theta)$
Mini-batch (size $B$ ):
$g_{\text{mini}} = \frac{1}{B}\sum_{i=1}^{B} \nabla_\theta \ell(x_{j_i}, y_{j_i}; \theta)$

Here the indices $j_i$ are randomly sampled examples.

Variance vs efficiency tradeoff:

Method	Gradient quality	Step cost	Steps per epoch	Typical use
Full-batch	Lowest variance, exact	Highest (all $N$ samples)	1	Small datasets, convex problems
Mini-batch (small B)	High variance, noisy	Very low	Many	Large-scale ML, exploration early in training
Mini-batch (larger B)	Lower variance	Higher	Fewer	Large LMs with GPUs/TPUs

Variance: With small batches, each batch’s gradient can point in a slightly different direction. Over steps, the noise partially cancels, approximating the full gradient.
Efficiency: Because each step is cheaper, you can take many more steps per unit time, often converging faster in wall-clock time than full-batch.

The noise in mini-batch gradients isn’t just a necessary evil; it also helps escape shallow local minima and poor plateaus, acting like a mild regularizer.

Concrete runnable example (PyTorch-style pseudo-code)

python

import torch
from torch import nn, optim

model = nn.Linear(10, 1)
optimizer = optim.SGD(model.parameters(), lr=0.1)
loss_fn = nn.MSELoss()

X = torch.randn(1000, 10)   # 1000 examples
y = torch.randn(1000, 1)

batch_size = 32
for epoch in range(5):
    perm = torch.randperm(X.size(0))
    X_shuf, y_shuf = X[perm], y[perm]
    for start in range(0, X.size(0), batch_size):
        xb = X_shuf[start:start + batch_size]
        yb = y_shuf[start:start + batch_size]

        pred = model(xb)
        loss = loss_fn(pred, yb)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

Each inner loop step uses a mini-batch of 32 samples.
If you changed batch_size = 1000, you’d be doing full-batch gradient descent.

1 / 4