So far, we’ve discussed different kinds of models and loss functions. We also mentioned that we need to tune the parameters $\theta$ of the model in some way that minimizes the loss function. Recall that for something simple like linear regression with MSE as the loss, we were able to find a best-solution formula, but usually, we won’t be so lucky. Since we aren’t able to find a best solution, we need some way to approximate it. In Deep Learning, the typical way to do so is by utilizing gradient descent (GD) methods. In this note, we’ll discuss different types of gradient descent methods.
These methods have two requirements:
- The collection of parameters should be a vector in some Euclidean space $\theta\in\mathbb R^M$.
- As a function of the parameter vector, the loss function $\mathbb R^M\xrightarrow{\theta\mapsto\ell(\mathscr D_\text{train}; \theta)}\mathbb R$ should be piecewise differentiable 1
Full-Batch Gradient Descent
The simplest variant of gradient descent operates on the entire training set at once. This is known as full-batch gradient descent. The procedure is as follows:
Algorithm parameters: learning rate $\eta > 0$, number of training steps $S$
Initialize $\theta$ by sampling from initialization distribution $\Theta$Loop for $t = 1, \dots, S$:
$\quad$ Compute gradient $g \leftarrow \nabla_\theta \ell(\mathscr{D}_\text{train}; \theta)$
$\quad$ Update $\theta \leftarrow \theta - \eta \cdot g$Return $\theta$
We begin by sampling $\theta$ from an initialization distribution $\Theta$. This choice of $\Theta$ matters quite a bit and is known as the initialization problem.
At each step, we compute the gradient $\nabla_\theta \ell(\mathscr{D}_\text{train}; \theta)$, which points in the direction of steepest loss increase. Moving opposite to it decreases the loss: $$ \theta \leftarrow \theta - \eta \cdot \nabla_\theta \ell(\mathscr{D}_\text{train}; \theta) $$ The scalar $\eta > 0$ is the learning rate, controlling the step size. It’s a hyperparameter, not a parameter of the model, but of the learning process (more on this in the in practice section below).
Training stops when some stopping condition is met. For now, we use the simplest one: a fixed number of steps, $S$ (our second hyperparameter).
Because full-batch GD uses the entire training set at every single step, each step here already amounts to one full pass over $\mathscr{D}_\text{train}$. Once we move to mini-batches, a full pass over the training set will take multiple steps, and this full pass is called an epoch.
Mini-Batch Gradient Descent
Instead of using the entire training set $\mathscr{D}_\text{train}$ at each step, we can use a small subset of it. Because we view the train set $\mathscr{D}_\text{train}$ as a collection of random samples from the joint distribution $(X, Y)$, we can simply extract small samples from this train set 2.
There are two variants: with replacement and without replacement.
With replacement
The idea here is that we may sample $\mathscr D'$ of size $N'$ with replacement from $\mathscr{D}_\text{train}$. Here, $N'$ is a new hyperparameter which we’ll call minibatch size 3. Then, we update the parameters as follows:
$$ \theta \leftarrow \theta - \eta \cdot \nabla_\theta \ell(\mathscr{D}'_\text{train}; \theta) $$
Note that we should always be taking a different small sample in each training step. Training on the same small sample can cause the model to overfit. The next variant of this aims to alleviate this problem.
Without replacement
This is the most commonly used version. We essentially perform mini-batch descent with random reshuffling. Because we sample without replacement, we ensure that all the data is taken into account. We still preserve the same hyperparameter $N'$ from the former variant, meaning that we’ll take $N'$-sized “slices” from $\mathscr{D}_\text{train}$. Note that we’ll end up with $\left\lfloor \frac{\left\lvert \mathscr{D}_{\text{train}} \right\rvert}{N'} \right\rfloor$ successive slices with the potential to have 1 more smaller slice. The choice of including this smaller slice comes with advantages and disadvantages: its inclusion will give you more data to work with, but if it’s too small then it may add too much noise. Because we have the choice of whether to include or discard this extra slice, it is yet another hyperparameter.
Putting this all together, the typical mini-batch GD algorithm is as follows:
Algorithm parameters: learning rate $\eta > 0$, minibatch size $N'$, number of slices per epoch $S'$
Initialize $\theta$ by sampling from initialization distribution $\Theta$Repeat (once per epoch):
$\quad$ Reshuffle $\mathscr{D}_\text{train} = \{(\underline{x}_i, y_i) : i \in \underline{I}\}$ according to a random permutation $\pi \in S_N$
$\quad$ Loop for $i = 0, \dots, S'-1$:
$\quad\quad$ Set $\mathscr{D}' \leftarrow \{(\underline{x}_{\pi(j)}, y_{\pi(j)}) : j = iN', \dots, \min((i+1)N',\, N) - 1\}$
$\quad\quad$ Compute gradient $g \leftarrow \nabla_\theta \ell(\mathscr{D}'; \theta)$
$\quad\quad$ Update $\theta \leftarrow \theta - \eta \cdot g$Return $\theta$
Each pass through the outer loop (one reshuffle followed by a full sweep through all $S'$ slices) is exactly one epoch. This is just a single pass over the entire training set, just split into $S'$ smaller updates instead of one big update. This is why you’ll often see training progress reported in epochs rather than raw step counts. It’s a size-independent way to talk about “how many times has the model seen the full dataset,” regardless of what minibatch size $N'$ you chose. Note that under this definition, one step of full-batch GD is also one epoch (it’s just the special case $S' = 1$).
But does random reshuffling truly outperform random sampling? Here’s a recent theoretical study addressing this question.
Stochastic Gradient Descent (SGD)
Stochastic4 gradient descent is, strictly speaking, the special case of mini-batch GD where $N' = 1$. Here, each update uses the gradient from a single training example.
A note on terminology:
In practice, this distinction gets blurred constantly. Deep learning papers, courses, and frameworks often say “SGD” when they actually mean mini-batch GD with $N' > 1$. For example, PyTorch’s optim.SGD accepts mini-batches of any size, not just $N' = 1$ (the “S” in its name refers to the family of methods). Also note that if you go look it up, you may also find other hyperparameters like momentum and weight decay, but don’t worry, since those are extensions to the base algorithm that we haven’t covered yet. Regardless, when you see “SGD” in the wild, make sure to check the context.
Here’s a brief summary of the three gradient descent variants we discussed:
| Method | $N'$ |
|---|---|
| SGD | $N' = 1$ |
| Mini-batch GD | $1 < N' < N$ |
| Full-batch GD | $N' = N$ (i.e. all of $\mathscr{D}_\text{train}$) |
Rather than just describing this in the abstract, it helps to see it: below is an interactive loss landscape you can experiment with directly.
The surface is $L(x, y) = \tfrac{1}{2}(x^2 + 8y^2)$. Try different learning rates and compare full-batch, mini-batch, and stochastic gradient descent. Observe that as $\eta$ increases, the steep direction transitions from smooth convergence to oscillation and eventually divergence.
In Practice
The theory above gives us a family of algorithms, but each one comes with knobs we have to set ourselves: $\eta$, $S$, $N'$, $S'$. These are chosen by us before training even starts, and in practice, getting good results depends a lot on choosing well.
Hyperparameter Tuning
Recall from above that a hyperparameter is a parameter of the learning process itself, as opposed to a parameter like $\theta$ that gets learned. Here, we have $\eta$, $S$, $N'$, and $S'$ as all the hyperparameters of gradient descent.
We could just try a bunch of values and keep whichever gives the lowest loss. But which loss? If we pick hyperparameters based on $\ell_\text{train}$, we run into a similar problem as overfitting $\theta$ itself: some hyperparameter settings (e.g. a very large $S$, or an $\eta$ that lets the model “memorize” $\mathscr{D}_\text{train}$) will drive training loss down while messing up generalization. And we already said $\mathscr{D}_\text{test}$ is off-limits until the very end. We can’t just use it to choose hyperparameters because that makes $\ell_\text{test}$ an unreliable estimate. So, we need a third dataset.
The Train-Validation-Test Split
We extend the train-test split into a train-validation-test split:
- A train set $\mathscr{D}_\text{train}$, used to fit $\theta$ for a given choice of hyperparameters.
- A validation set $\mathscr{D}_\text{val}$, used to compare different hyperparameter choices against each other.
- A test set $\mathscr{D}_\text{test}$, touched exactly once, at the very end, to report a final unbiased estimate of generalization.
A typical split ratio might be 70/15/15 or 80/10/10, partitioned randomly (with a fixed seed, for reproducibility) just like before.
The final workflow looks like this: for each candidate setting of $(\eta, S, N', S')$, train a model on $\mathscr{D}_\text{train}$ and evaluate it on $\mathscr{D}_\text{val}$. Pick the setting with the lowest $\ell_\text{val}$. Only then, with your hyperparameters and $\theta$ both fixed, do you compute $\ell_\text{test}$ once to report how the model will actually perform on new data.
In this process, $\mathscr{D}_\text{val}$ is allowed to influence your choice of hyperparameters. We’ll tune the hyperparameters such that $\mathscr{D}_\text{val}$ is good enough, but by the time you touch $\mathscr{D}_\text{test}$, every decision must already been locked in. This preserves the “integrity” of $\mathscr{D}_\text{test}$, which is why that’s the number we report at the end.
How you actually search over hyperparameter settings (like trying every combination on a fixed grid, or sampling combinations at random) is its own design choice, with its own tradeoffs. This is a topic we may cover in future note.
-
This means that the function can be split up into sub-domains where it each “piece” is differentiable, but the function in its entirety may contain sharp corners or breaks where the pieces meet. ↩
-
This is quite similar to bootstrapping (if you happen to be familiar); we simulate sampling from the distribution $(X, Y)$ by sampling from the training set $\mathscr{D}_\text{train}$. ↩
-
Conventionally, minibatch size $N'=2^{n'}$ where $n' \in \mathbb{N}$ is any number; however, this convention is being challenged. ↩
-
Stochastic means random. In SGD, the gradient at each step is computed from randomly selected training data instead of the entire dataset. ↩