From 99e5781c911ed93f524d8c166e6b2a5d7d183689 Mon Sep 17 00:00:00 2001 From: Rick McEwen Date: Mon, 5 Jan 2026 11:01:14 -0500 Subject: [PATCH] Fix trainer to use separation as sole criterion for best model Previously the trainer saved a new "best" model if either separation OR loss improved, with loss checked as a fallback. This caused confusing behavior where models with lower separation could overwrite better models. Now only separation (gap between positive and negative similarity) is used to determine the best model, which is the key metric for contrastive learning quality. --- training/trainer.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/training/trainer.py b/training/trainer.py index e624283..07a8da6 100644 --- a/training/trainer.py +++ b/training/trainer.py @@ -169,16 +169,11 @@ class Trainer: "val_neg_sim": val_metrics["mean_neg_sim"], }) - # Checkpointing based on separation (primary) or loss (secondary) - improved = False + # Checkpointing based on separation (gap between pos and neg similarity) + # This is the key metric for contrastive learning quality if val_metrics["separation"] > self.best_val_separation + self.config.min_delta: self.best_val_separation = val_metrics["separation"] - improved = True - elif val_metrics["loss"] < self.best_val_loss - self.config.min_delta: - self.best_val_loss = val_metrics["loss"] - improved = True - - if improved: + self.best_val_loss = val_metrics["loss"] # Track for reference self.patience_counter = 0 self._save_checkpoint("best.pt") self.logger.info("New best model saved!")