diff --git a/crates/ml-alpha/examples/alpha_train.rs b/crates/ml-alpha/examples/alpha_train.rs index 6f20d31b5..f3705b8cf 100644 --- a/crates/ml-alpha/examples/alpha_train.rs +++ b/crates/ml-alpha/examples/alpha_train.rs @@ -82,10 +82,19 @@ struct Cli { lr_min_factor: f32, /// Early-stop training after this many consecutive epochs without - /// val_loss improvement. 0 disables early stopping. - #[arg(long, default_value_t = 3)] + /// improvement on `early_stop_metric`. 0 disables early stopping. + #[arg(long, default_value_t = 5)] early_stop_patience: usize, + /// Which metric drives early stopping: "val_loss" (stable, tracks + /// probability calibration) or "mean_auc" (noisier, tracks the + /// ranking quality used for downstream trading). mean_auc needs + /// higher patience (~5) because it bounces 1-2pt epoch-to-epoch + /// even when the model is still genuinely improving on the long + /// horizons that the auto-horizon-weights down-weight. + #[arg(long, default_value = "mean_auc")] + early_stop_metric: String, + /// Custom per-horizon BCE weights (5 floats, comma-separated). When /// set, overrides the auto/uniform defaults. Example: "1.0,1.0,0.5,0.1,0.02". #[arg(long)] @@ -206,11 +215,20 @@ fn main() -> Result<()> { let mut final_val_auc = [0.5_f32; N_HORIZONS]; let mut n_val_seqs_consumed = 0usize; + // Validate early-stop metric choice upfront. + let early_stop_metric = cli.early_stop_metric.as_str(); + anyhow::ensure!( + matches!(early_stop_metric, "val_loss" | "mean_auc" | "none"), + "--early-stop-metric must be one of: val_loss, mean_auc, none (got `{}`)", + early_stop_metric + ); + // Best-checkpoint tracking (by val_loss). let mut best_val_loss = f32::INFINITY; let mut best_epoch = 0usize; let mut best_val_auc = [0.5_f32; N_HORIZONS]; - let mut epochs_since_improvement = 0usize; + let mut val_loss_no_improvement = 0usize; + let mut mean_auc_no_improvement = 0usize; let mut early_stopped = false; let mut epochs_completed = 0usize; @@ -370,40 +388,47 @@ fn main() -> Result<()> { n_val_seqs_consumed = val_loader.yielded(); epochs_completed = epoch + 1; - // Best-checkpoint + early-stop bookkeeping (by val_loss). - if val_avg < best_val_loss { + // Best-checkpoint trackers — independent for each metric. + let val_loss_improved = val_avg < best_val_loss; + if val_loss_improved { best_val_loss = val_avg; best_epoch = epoch; best_val_auc = per_horizon_auc; - epochs_since_improvement = 0; - tracing::info!( - epoch, val_loss = val_avg, "new best val_loss" - ); + val_loss_no_improvement = 0; + tracing::info!(epoch, val_loss = val_avg, "new best val_loss"); } else { - epochs_since_improvement += 1; - if cli.early_stop_patience > 0 && epochs_since_improvement >= cli.early_stop_patience { + val_loss_no_improvement += 1; + } + + let mean_auc = per_horizon_auc.iter().sum::() / per_horizon_auc.len() as f32; + let mean_auc_improved = mean_auc > best_mean_auc; + if mean_auc_improved { + best_mean_auc = mean_auc; + best_mean_auc_epoch = epoch; + best_mean_auc_per_horizon = per_horizon_auc; + mean_auc_no_improvement = 0; + tracing::info!(epoch, mean_auc, "new best mean_auc"); + } else { + mean_auc_no_improvement += 1; + } + + // Early-stop decision based on selected metric. + if cli.early_stop_patience > 0 { + let (counter, metric_label, best_val) = match early_stop_metric { + "val_loss" => (val_loss_no_improvement, "val_loss", best_val_loss), + "mean_auc" => (mean_auc_no_improvement, "mean_auc", best_mean_auc), + _ => (0, "none", 0.0), // disabled + }; + if counter >= cli.early_stop_patience && metric_label != "none" { tracing::warn!( - epoch, best_epoch, best_val_loss, + epoch, metric = metric_label, best = best_val, patience = cli.early_stop_patience, - "early stopping — val_loss did not improve", + "early stopping", ); early_stopped = true; break; } } - - // Independent best-by-mean-AUC tracker — reported only, does NOT - // gate early stopping (val_loss remains the early-stop signal so - // the two policies stay decoupled). - let mean_auc = per_horizon_auc.iter().sum::() / per_horizon_auc.len() as f32; - if mean_auc > best_mean_auc { - best_mean_auc = mean_auc; - best_mean_auc_epoch = epoch; - best_mean_auc_per_horizon = per_horizon_auc; - tracing::info!( - epoch, mean_auc, "new best mean_auc" - ); - } } let summary = AlphaTrainSummary { diff --git a/infra/k8s/argo/alpha-perception-template.yaml b/infra/k8s/argo/alpha-perception-template.yaml index f9d487a68..4fb499b44 100644 --- a/infra/k8s/argo/alpha-perception-template.yaml +++ b/infra/k8s/argo/alpha-perception-template.yaml @@ -66,6 +66,10 @@ spec: value: "1" - name: auto-horizon-weights value: "false" + - name: early-stop-metric + value: "mean_auc" + - name: early-stop-patience + value: "5" volumes: - name: git-ssh-key @@ -431,6 +435,8 @@ spec: --n-val-seqs {{workflow.parameters.n-val-seqs}} \ --seed {{workflow.parameters.seed}} \ --batch-size {{workflow.parameters.batch-size}} \ + --early-stop-metric {{workflow.parameters.early-stop-metric}} \ + --early-stop-patience {{workflow.parameters.early-stop-patience}} \ $EXTRA_FLAGS echo "=== Training complete ===" diff --git a/scripts/argo-alpha-perception.sh b/scripts/argo-alpha-perception.sh index 0e9c5b7b1..8f8cad0de 100755 --- a/scripts/argo-alpha-perception.sh +++ b/scripts/argo-alpha-perception.sh @@ -29,6 +29,8 @@ N_VAL_SEQS=1000 SEED=16962 BATCH_SIZE=1 AUTO_HORIZON_WEIGHTS=false +EARLY_STOP_METRIC=mean_auc +EARLY_STOP_PATIENCE=5 WATCH=false usage() { @@ -64,6 +66,8 @@ while [[ $# -gt 0 ]]; do --seed) SEED="$2"; shift 2 ;; --batch-size) BATCH_SIZE="$2"; shift 2 ;; --auto-horizon-weights) AUTO_HORIZON_WEIGHTS=true; shift ;; + --early-stop-metric) EARLY_STOP_METRIC="$2"; shift 2 ;; + --early-stop-patience) EARLY_STOP_PATIENCE="$2"; shift 2 ;; --watch) WATCH=true; shift ;; -h|--help) usage; exit 0 ;; *) echo "Unknown option: $1"; usage; exit 1 ;; @@ -129,4 +133,6 @@ argo submit -n foxhunt --from=wftmpl/alpha-perception \ -p seed="$SEED" \ -p batch-size="$BATCH_SIZE" \ -p auto-horizon-weights="$AUTO_HORIZON_WEIGHTS" \ + -p early-stop-metric="$EARLY_STOP_METRIC" \ + -p early-stop-patience="$EARLY_STOP_PATIENCE" \ $WATCH_FLAG