Checkpointing And Early Stopping

Why this matters

In real NLP projects, training is expensive and fragile. Checkpointing protects progress (so you can resume or roll back) and early stopping prevents overfitting while saving compute. You will use both when fine-tuning language models, running long experiments on shared GPUs, or shipping production models where reliability matters.

• Fine-tuning a sentiment classifier and resuming after a preemption.
• Stopping summarization training when ROUGE stops improving to save GPU hours.
• Keeping the best NER model snapshot for deployment and reproducibility.

Concept explained simply

Checkpointing: periodically save your model state (weights + training context) so you can resume training or deploy the best version later.

Early stopping: monitor a validation metric and stop when it stops improving for a while. This avoids overfitting and wasted time.

Mental model

Imagine hiking up a mountain:

• Checkpoints are safe camps you set along the trail. If weather turns bad, you can return to the last safe camp.
• Early stopping is your rule to turn back if the view (validation metric) hasn’t improved after several lookouts (patience).

Key components

• What to save in a checkpoint:
  • Model weights (mandatory)
  • Optimizer state (for proper resume)
  • LR scheduler state (if used)
  • Epoch/step number and best metric so far
  • Random seeds (optional but helpful)
  • NLP specifics: tokenizer/vocabulary, label mapping, preprocessing configs
• When to save:
  • Every N steps/epochs (periodic)
  • On improvement of a monitored metric (best-only or top-k)
  • On time-based intervals (long runs)
• Retention policy:
  • Keep only the best k checkpoints
  • Budget-based cleanup to manage disk usage
• Early stopping knobs:
  • metric: e.g., val_loss (minimize), F1/ROUGE/BLEU/accuracy (maximize)
  • mode: "min" or "max"
  • min_delta: minimum meaningful change to count as improvement
  • patience: allowed number of checks with no improvement
  • cooldown: optional wait before resuming monitoring after an improvement
  • restore_best_weights: load the best checkpoint at the end
• Noise handling:
  • Use moving average or evaluate less frequently to reduce metric noise
  • Use a sensible min_delta to ignore tiny fluctuations
• Reproducibility:
  • Save configs and seeds with the checkpoint
  • Record library versions

Worked examples

How to implement (step-by-step)

# EarlyStopping helper
class EarlyStopper:
    def __init__(self, mode='min', patience=3, min_delta=0.0, restore_best=True):
        self.mode = mode
        self.patience = patience
        self.min_delta = min_delta
        self.best = None
        self.num_bad = 0
        self.restore_best = restore_best
        self.best_state = None

    def is_better(self, current):
        if self.best is None:
            return True
        if self.mode == 'min':
            return current < self.best - self.min_delta
        else:
            return current > self.best + self.min_delta

    def step(self, metric, model_state=None):
        if self.is_better(metric):
            self.best = metric
            self.num_bad = 0
            if model_state is not None:
                self.best_state = {k: v.cpu().clone() for k, v in model_state.items()}
            return True  # improved
        else:
            self.num_bad += 1
            return False  # no improvement

    def should_stop(self):
        return self.num_bad >= self.patience

# Training loop (outline)
best_path = 'best.pt'
for epoch in range(num_epochs):
    train_one_epoch(model, optimizer, train_loader)
    val_metric = evaluate(model, val_loader)  # e.g., val_loss (lower is better)

    improved = stopper.step(val_metric, model.state_dict())
    if improved:
        save_checkpoint({
            'model': model.state_dict(),
            'optimizer': optimizer.state_dict(),
            'scheduler': scheduler.state_dict() if scheduler else None,
            'epoch': epoch,
            'best_metric': stopper.best,
            'tokenizer': tokenizer_state,  # serialize vocab if needed
            'config': train_config,
        }, path=best_path)

    if stopper.should_stop():
        break

if stopper.restore_best and stopper.best_state is not None:
    model.load_state_dict(stopper.best_state)

from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint

callbacks = [
    ModelCheckpoint(
        filepath='best.h5',
        monitor='val_loss',
        mode='min',
        save_best_only=True,
        save_weights_only=True,
        verbose=1,
    ),
    EarlyStopping(
        monitor='val_loss',
        mode='min',
        patience=3,
        min_delta=0.002,
        restore_best_weights=True,
        verbose=1,
    )
]

model.fit(train_ds, validation_data=val_ds, epochs=20, callbacks=callbacks)

Exercises

Note: The quick test is available to everyone; only logged-in users get saved progress.

Common mistakes and self-check

• Monitoring training loss instead of validation: use a held-out set.
• min_delta too small: tiny fluctuations count as improvements; increase it.
• Wrong mode: maximizing a loss or minimizing an accuracy. Double-check.
• Saving only weights: cannot resume optimizer/scheduler states.
• No retention policy: disk fills up; set top-k or cleanup.
• Not restoring best weights: you might deploy a worse epoch.
• Evaluating too often with noisy metrics: smooth or evaluate less frequently.
• Tuning on the test set: creates leakage; keep a final test untouched.

Practical projects

• Fine-tune a small BERT on sentiment. Compare runs with and without early stopping. Record GPU hours saved.
• Simulate a preemption: kill training mid-epoch and resume from the latest checkpoint. Verify no metric regressions.
• Top-k policy: keep best 3 checkpoints by F1 on NER. After training, ensemble the top-3 and compare to single best.

Learning path

• Before this: Dataset splitting and metrics, Optimizers/LR schedulers, Reproducibility basics.
• Now: Checkpointing and Early Stopping (this page).
• Next: Learning rate scheduling strategies, Mixed precision and gradient accumulation, Automated hyperparameter search with early stopping.

Who this is for

• NLP engineers fine-tuning transformer models
• ML engineers running long GPU jobs
• Data scientists preparing robust experiments

Prerequisites

• Python and basic deep learning
• Intro knowledge of PyTorch or Keras/TensorFlow
• Understanding of validation metrics and splits

Mini challenge

Your validation F1 oscillates: 89.8, 90.0, 90.1, 90.05, 90.08, 90.09, 90.07. Choose mode, min_delta, and patience to balance stability and speed. Explain your choice in 2–3 sentences. Then pick a retention policy for checkpoints.

Next steps

• Implement early stopping in your current project; log each decision (improved/no-improve).
• Switch from periodic to best-only checkpoints and measure disk and time savings.
• Move on to the quick test to check your understanding.