From 1f2c91f2e580ecd84a7ceb8d6707653f2b642a78 Mon Sep 17 00:00:00 2001
From: stephantul <stephantul@gmail.com>
Date: Mon, 2 Nov 2026 21:07:01 +0100
Subject: [PATCH 1/2] fix: scheduler should step more often

---
 model2vec/train/trainer.py | 17 +++--------------
 tests/test_trainable.py    |  6 +++---
 2 files changed, 6 insertions(+), 17 deletions(-)

diff --git a/model2vec/train/trainer.py b/model2vec/train/trainer.py
index dea7192..c7783e5 100644
--- a/model2vec/train/trainer.py
+++ b/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
diff --git a/tests/test_trainable.py b/tests/test_trainable.py
index a8a4b21..8160cb5 100644
--- a/tests/test_trainable.py
+++ b/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:

From 334f2ea1a3f13b3fd3e86680457a994c1844003e Mon Sep 17 00:00:00 2001
From: stephantul <stephantul@gmail.com>
Date: Mon, 2 Nov 2026 21:14:51 +0100
Subject: [PATCH 2/2] address comment

---
 model2vec/train/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/model2vec/train/README.md b/model2vec/train/README.md
index 3596cb2..b07b3bb 100644
--- a/model2vec/train/README.md
+++ b/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.