Add early stopping patience to reduce variance

adamkarvonen · Oct 25, 2024 · 5285563 · 5285563
1 parent d5a215a
commit 5285563
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/evals/shift_and_tpp/eval_config.py b/evals/shift_and_tpp/eval_config.py
@@ -14,6 +14,10 @@ class EvalConfig:
     )
 
     spurious_corr: bool = True
+
+    # This reduces randomness in the SCR results
+    early_stopping_patience: int = 40
+
     # Load datset and probes
     train_set_size: int = 4000
     test_set_size: int = 1000  # This is limited as the test set is smaller than the train set

diff --git a/evals/shift_and_tpp/main.py b/evals/shift_and_tpp/main.py
@@ -521,6 +521,7 @@ def run_eval_single_dataset(
             epochs=config.probe_epochs,
             lr=config.probe_lr,
             spurious_corr=config.spurious_corr,
+            early_stopping_patience=config.early_stopping_patience,
         )
 
         torch.set_grad_enabled(False)