In the last note, we introduced ensembles: training $n$ independently-seeded copies of a model at once, both to get confidence intervals on performance and, typically, better predictions from aggregating over the ensemble. This note builds directly on that setup. Weâll cover a way to keep the best-scoring checkpoint of each ensemble member during training, a stopping condition built on top of that, and masking, an indexing trick that makes this implementable without looping over the ensemble. We also briefly discuss grid search, which is probably the simplest way to search over hyperparameter settings (which finally answers a question left open in the gradient descent note). This note is more practical-flavored than most of the others.
Selecting the Best Model
A modelâs validation score doesnât necessarily keep improving for the whole training run. Itâs actually quite common for it to peak and then get worse, and this is usually due to the model overfitting (we may discuss more on this later). Due to this, the model we get when training stops might not be the best version of the model. So, we need some sort of way to pluck out the version of this model that had the best validation score along the way.
Suppose weâre training an ensemble of models rather than a single one (and itâs usually good practice to do so if possible). Let $E$ be the set of indices into the ensemble. If the ensemble is shaped (n_1,...,n_l), then $E = [n_1]\times\dotsb\times[n_l]$, i.e., one index per ensemble member. Weâll assume higher scores are better (accuracy, say); but if our metric is the opposite (like a loss), we can just flip the inequalities below.
The âbookkeepingâ works like this:
- Setup:
- Keep a running best score $s'_e$ for each ensemble member $e \in E$, initialized to something no better than the worst possible score (e.g. $0$ for accuracy).
- Keep a running best set of parameters $\theta'_e$ for each $e$, initialized however since theyâll get overwritten on the first validation pass anyway.
- Every time you validate:
- You get a fresh score $s_e$ and the current parameters $\theta_e$ for each ensemble member.
- For any member whose new score beats its running best, update both the best score and the best parameters. Because floating-point comparisons are noisy, âbeatsâ means beats by more than some small improvement threshold $\delta>0$, not just $s_e > s'_e$:
$$ s'_e\leftarrow s_e\text{ and }\theta'_e\leftarrow\theta_e \text{ for }e\in E\text{ such that }s_e > s'_e + \delta $$
This update happens independently member-by-member where one ensemble member plateauing doesnât stop the others from continuing to record new bests.
Early Stopping
Once youâre tracking best scores this way, you get a natural stopping rule: if itâs been a while since any ensemble member last improved, further training probably isnât worth it. So weâll just stop once youâve gone $C$ training steps without an improvement. $C$ is usually called the patience (new hyperparameter!)
Concretely, extend the bookkeeping above with a counter:
- Setup: also keep a counter $c$ of training steps since the last improvement, starting at $0$.
- Every time you validate:
- If at least one ensemble member improved, reset $c$ to $0$.
- Otherwise, add the number of training steps since the last validation to $c$, and if $c \ge C$, stop training.
Note that since $c$ is shared across the whole ensemble, as long as some member is still improving, training continues. The run only stops once every member has simultaneously stalled for $C$ steps in a row.
Masking
Thereâs still a gap between the update rule above and actually implementing it: how do you apply an update to only the ensemble members that improved, without writing an explicit loop over $E$? The answer is masking. Weâre basically indexing into a tensor using another tensor of True/False values instead of integer indices.
Itâs easiest to see with a small example before getting formal about it. Suppose we have a 1-D tensor of scores and a same-shaped tensor telling us which ones improved:
scores = torch.tensor([0.81, 0.77, 0.93, 0.60])
improved = torch.tensor([True, False, True, False])
scores[improved] # -> tensor([0.81, 0.93])
scores[improved] keeps only the entries where the mask is True, in order, and filters out the rest. Two things about this generalize to what we actually need:
- The mask doesnât have to cover the tensorâs entire shape. If each ensemble memberâs parameters are themselves a vector (or some higher-dimensional tensor), a mask over just the ensemble dimensions will select whole parameter vectors at once, not individual scalars.
- Masking isnât only for reading but assigning as well.
best_scores[improved] = scores[improved]writes the new scores into exactly the entries that improved and leaves everything else alone. Thatâs the update rule from the previous section, written as one line of code instead of a loop.
Here's the general, more precise version of this, stated precisely.
I acknowledge that this formalization might have resulted in a difficult presentation, so just consider this as something Iâm obligated to do as a responsible holder of a bachelorâs degree in math.
Let $[n_1]\times\dotsb\times[n_k]\xrightarrow tK$ be an arbitrary tensor, and let $[n_1]\times\dotsb\times[n_l]\xrightarrow b\{\mathtt{True}, \mathtt{False}\}$ be a Boolean tensor over some prefix of $t$âs dimensions ($l\le k$). Collect the coordinates where $b$ is True into a list $\mathscr I=(I_j:0\le j<r)$, with $I_j=(i_1^{(j)},\dotsc,i_l^{(j)})$. Then $a[b]$ is defined as the composite
$$ [r]\times[n_{l+1}]\times\dotsb\times[n_k] \xrightarrow{(j,i_{l+1},\dotsc,i_k)\mapsto(i_1^{(j)},\dotsc,i_l^{(j)},i_{l+1},\dotsc,i_k)} [n_1]\times\dotsb\times[n_k]\xrightarrow aK. $$
In words: line up the coordinates where $b$ is True, and stack the corresponding subtensors of $a$ along a new first axis of length $r$. Array libraries that support this kind of indexing, like NumPy, compute $r$ and $\mathscr I$ for you internally, but youâll rarely need either directly since indexing directly with the mask handles both reading and writing on its own.
We can tie this back to the best model selector. Let $b$ be the Boolean tensor $e\mapsto[s_e > s'_e + \delta]$ over $E$. Then
$$ s'[b]\leftarrow s[b]\text{ and }\theta'[b]\leftarrow\theta[b] $$
is the update applied to only the ensemble members that improved.
For more information, see NumPyâs documentation on Boolean array indexing.
Grid Search
Recall from the gradient descent note that we still owed an answer for how to actually search over candidate hyperparameter settings, evaluating each on $\mathscr D_\text{val}$. The simplest hyperparameter search strategy is Grid Search: pick a finite set of candidate values for each hyperparameter you care about, then just try every combination.
If hyperparameter $i$ has $m_i$ candidate values, thatâs $\prod_i m_i$ configurations total. Note that this means that the count multiplies across hyperparameters, which is why grid search gets expensive fast as you add more of them.
Thatâs it.
In Practice
- After each validation pass, compute the improvement mask, use it to update
best_scoresandbest_state_dict(or a stacked parameter tensor if youâre running an ensemble), and update the patience counter. Itâs may be wise to periodically writebest_state_dictto disk rather than only at the very end of training, so a crash doesnât cost you the best checkpoint. - The threshold comparison is just
improved = scores > best_scores + delta, and the write-back isbest_scores[improved] = scores[improved](same pattern for the parameters). - A minimal early-stopping check:
if improved.any(): counter = 0 else: counter += steps_since_last_validation stop = counter >= patience - For grid search,
itertools.productover your candidate lists is usually enough to enumerate all $\prod_i m_i$ configurations by hand. Once the grid gets large, dedicated hyperparameter search libraries (e.g. scikit-learnâsGridSearchCV, orray.tune) will run the loop for you, and some offer smarter search strategies beyond a plain grid searching. - Grid search spends just as much effort on clearly bad regions of the hyperparameter space as on promising ones. Random search, and later, adaptive methods like Bayesian optimization, tend to scale better as you add more hyperparameters.