fix: scheduler should step more often

model2vec/train/README.md

@@ -227,7 +227,7 @@ Without an explicit validation set, `fit` holds out `test_size` of the data for
 
 By default, the model is validated once per epoch. If an epoch has more than 250 batches, it is instead validated about four times per epoch, but at most once every 250 steps. `validation_steps` sets a fixed number of steps between validations instead.
 
-Classifiers stop early on the validation accuracy, the other models on the validation loss. Whenever the validation loss stops improving for a few epochs, the learning rate is halved. After training, the weights from the best validation check are loaded back into the model.
+Classifiers stop early on the validation accuracy, the other models on the validation loss. Whenever the validation loss fails to improve by at least 3% for three validation checks in a row, the learning rate is halved, down to a minimum of `1e-6`. This counts validation checks, not epochs, so with several checks per epoch the learning rate can drop within a single epoch, and a smaller `validation_steps` makes it drop sooner. After training, the weights from the best validation check are loaded back into the model.
 
 The training loop shows the loss, and for classifiers the accuracy, of the latest training batches and of the latest validation check in the progress bar.
 

model2vec/train/trainer.py

@@ -47,14 +47,6 @@ def _resolve_max_epochs(max_epochs: int | None) -> int:
     return max_epochs
 
 
-def _step_plateau_scheduler(
-    scheduler: torch.optim.lr_scheduler.ReduceLROnPlateau, latest_val_loss: float | None
-) -> None:
-    """Step a ReduceLROnPlateau scheduler with the latest val_loss, once per epoch, if one is available yet."""
-    if latest_val_loss is not None:
-        scheduler.step(latest_val_loss)
-
-
 def resolve_device(device: str) -> torch.device:
     """Resolve a device string, with "auto" picking the best available accelerator."""
     if device != "auto":
@@ -132,7 +124,7 @@ def run_training_loop(  # noqa: C901
         optimizer,
         mode="min",
         factor=0.5,
-        patience=3,
+        patience=2,
         min_lr=1e-6,
         threshold=0.03,
         threshold_mode="rel",
@@ -151,12 +143,11 @@ def run_training_loop(  # noqa: C901
     global_step = 0
     postfix: dict[str, str] = {}
     train_metric_windows: dict[str, deque[float]] = defaultdict(lambda: deque(maxlen=_TRAIN_METRICS_WINDOW))
-    latest_val_loss: float | None = None
 
     def validate_and_checkpoint() -> bool:
-        nonlocal best_checkpoint, best_val_metric, latest_val_loss
+        nonlocal best_checkpoint, best_val_metric
         metrics = _run_validation(model, loss_function, compute_metrics, val_loader, device)
-        latest_val_loss = metrics["val_loss"]
+        scheduler.step(metrics["val_loss"])
         current = metrics[val_metric]
         improved = current < best_val_metric if early_stopping_direction == "min" else current > best_val_metric
         if improved:
@@ -200,7 +191,5 @@ def validate_and_checkpoint() -> bool:
                 if should_stop and (min_epochs is None or current_epoch >= min_epochs):
                     return best_checkpoint
 
-            _step_plateau_scheduler(scheduler, latest_val_loss)
-
             if current_epoch >= max_epochs:
                 return best_checkpoint

tests/test_trainable.py

@@ -1196,8 +1196,8 @@ def test_run_training_loop_mid_epoch_early_stop() -> None:
     assert set(state_dict) == set(model.state_dict())
 
 
-def test_run_training_loop_steps_scheduler_once_per_epoch(monkeypatch: pytest.MonkeyPatch) -> None:
-    """The LR scheduler steps once per epoch, not once per validation check."""
+def test_run_training_loop_steps_scheduler_once_per_validation_check(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The LR scheduler steps once per validation check, not once per epoch."""
     step_calls: list[float] = []
     original_step = torch.optim.lr_scheduler.ReduceLROnPlateau.step
 
@@ -1223,7 +1223,7 @@ def counting_step(self: torch.optim.lr_scheduler.ReduceLROnPlateau, metrics: flo
         val_check_interval=1,
         check_val_every_epoch=None,
     )
-    assert len(step_calls) == 3
+    assert len(step_calls) == 18
 
 
 def test_split_indices() -> None: