Early stopping: monitor validation loss, restore the best checkpoint
The problem, from first principles
Overparameterization and double descent: more parameters than data can still generalize, later in this Part, complicates a simpler, older intuition worth understanding first: train a model long enough on a fixed training set, and TRAINING loss will keep decreasing (the model can always keep fitting its training examples a little better), but VALIDATION loss, measured on data the model never trains on, often stops improving and starts getting WORSE partway through training, the model has begun memorizing training-set-specific noise rather than learning anything that generalizes. The training run itself doesn't know this is happening; left running for its full planned number of epochs, it'll happily keep optimizing training loss long past the point where validation performance peaked.
Early stopping is the direct fix: watch validation loss every epoch, and if it hasn't improved for a while (a "patience" window of consecutive epochs), stop training and go back to whichever checkpoint actually had the BEST validation loss, not the checkpoint from whenever training happened to end.
From theory to code
Implement EarlyStopping.step(val_loss, state), called once per epoch. If val_loss is better (lower) than the best ever seen by more than min_delta, record it as the new best (along with state, typically a snapshot of the model's weights at that point) and reset the patience counter to zero. Otherwise, increment the patience counter, and if it reaches self.patience, set self.should_stop = True. Return self.should_stop either way, so a training loop can check the return value directly.
Constraints
- The very FIRST call to
stepalways counts as an improvement (there's no previous best to compare against yet), regardless of how largeval_lossis. - An improvement must beat the previous best by MORE than
min_delta, not just be strictly smaller:val_loss < best_loss - min_delta, notval_loss < best_loss. Withmin_delta=0, this reduces to any strict improvement counting. self.counterresets to0on every improvement, and only accumulates on CONSECUTIVE non-improving epochs; a single improvement in the middle of a losing streak resets the streak entirely.self.best_statemust be updated to whateverstatewas passed on the epoch that produced the new bestval_loss, so it always reflects the actual best-performing checkpoint.
Hints
Hint 1: The first call
self.best_loss is None is only true before the first call, use if self.best_loss is None or val_loss < self.best_loss - self.min_delta: as the single condition covering both "this is the first epoch" and "this epoch genuinely improved."
Hint 2: Improvement vs. no improvement
Inside that if, update self.best_loss = val_loss, self.best_state = state, and reset self.counter = 0. In the else branch, increment self.counter += 1, and if self.counter >= self.patience, set self.should_stop = True.
Hint 3
Return self.should_stop as the last line, this lets a training loop write if early_stopping.step(val_loss): break directly, without needing to check a separate attribute afterward.
Click "Run Tests" to test your implementation