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.
This commit is contained in:
Rick McEwen
2026-01-05 11:01:14 -05:00
parent 44e8b6ae7d
commit 99e5781c91

View File

@ -169,16 +169,11 @@ class Trainer:
"val_neg_sim": val_metrics["mean_neg_sim"], "val_neg_sim": val_metrics["mean_neg_sim"],
}) })
# Checkpointing based on separation (primary) or loss (secondary) # Checkpointing based on separation (gap between pos and neg similarity)
improved = False # This is the key metric for contrastive learning quality
if val_metrics["separation"] > self.best_val_separation + self.config.min_delta: if val_metrics["separation"] > self.best_val_separation + self.config.min_delta:
self.best_val_separation = val_metrics["separation"] self.best_val_separation = val_metrics["separation"]
improved = True self.best_val_loss = val_metrics["loss"] # Track for reference
elif val_metrics["loss"] < self.best_val_loss - self.config.min_delta:
self.best_val_loss = val_metrics["loss"]
improved = True
if improved:
self.patience_counter = 0 self.patience_counter = 0
self._save_checkpoint("best.pt") self._save_checkpoint("best.pt")
self.logger.info("New best model saved!") self.logger.info("New best model saved!")