diff --git a/crates/ml/examples/hyperopt_baseline_rl.rs b/crates/ml/examples/hyperopt_baseline_rl.rs index 9ddfac974..4ef6340f3 100644 --- a/crates/ml/examples/hyperopt_baseline_rl.rs +++ b/crates/ml/examples/hyperopt_baseline_rl.rs @@ -169,13 +169,14 @@ fn run_dqn_hyperopt(args: &Args, parallel: usize, gpu_devices: &[candle_core::De let start = Instant::now(); let result = if args.optimizer.as_str() == "tpe" { info!("Using TPE (Tree-Parzen Estimator) optimizer"); - ml::hyperopt::optimize_with_tpe(trainer, args.trials, args.n_initial, Some(args.seed)) + ml::hyperopt::optimize_with_tpe(trainer, args.trials, args.n_initial, Some(args.seed), Some("dqn")) .context("DQN TPE hyperopt optimization failed")? } else { let optimizer = ArgminOptimizer::builder() .max_trials(args.trials) .n_initial(args.n_initial) .seed(args.seed) + .model_name("dqn") .build(); if parallel > 1 { @@ -249,13 +250,14 @@ fn run_ppo_hyperopt(args: &Args, parallel: usize, gpu_devices: &[candle_core::De let start = Instant::now(); let result = if args.optimizer.as_str() == "tpe" { info!("Using TPE (Tree-Parzen Estimator) optimizer"); - ml::hyperopt::optimize_with_tpe(trainer, args.trials, args.n_initial, Some(args.seed)) + ml::hyperopt::optimize_with_tpe(trainer, args.trials, args.n_initial, Some(args.seed), Some("ppo")) .context("PPO TPE hyperopt optimization failed")? } else { let optimizer = ArgminOptimizer::builder() .max_trials(args.trials) .n_initial(args.n_initial) .seed(args.seed) + .model_name("ppo") .build(); if parallel > 1 { diff --git a/crates/ml/examples/hyperopt_baseline_supervised.rs b/crates/ml/examples/hyperopt_baseline_supervised.rs index 7c576d2dd..9a1e4afc6 100644 --- a/crates/ml/examples/hyperopt_baseline_supervised.rs +++ b/crates/ml/examples/hyperopt_baseline_supervised.rs @@ -136,6 +136,7 @@ fn run_tft_hyperopt(args: &Args) -> Result { .max_trials(args.trials) .n_initial(args.n_initial) .seed(args.seed) + .model_name("tft") .build(); training_metrics::set_hyperopt_trial("tft", 0.0, args.trials as f64); @@ -195,6 +196,7 @@ fn run_mamba2_hyperopt(args: &Args) -> Result { .max_trials(args.trials) .n_initial(args.n_initial) .seed(args.seed) + .model_name("mamba2") .build(); training_metrics::set_hyperopt_trial("mamba2", 0.0, args.trials as f64); @@ -255,6 +257,7 @@ fn run_liquid_hyperopt(args: &Args) -> Result { .max_trials(args.trials) .n_initial(args.n_initial) .seed(args.seed) + .model_name("liquid") .build(); training_metrics::set_hyperopt_trial("liquid", 0.0, args.trials as f64); @@ -315,6 +318,7 @@ fn run_tggn_hyperopt(args: &Args) -> Result { .max_trials(args.trials) .n_initial(args.n_initial) .seed(args.seed) + .model_name("tggn") .build(); training_metrics::set_hyperopt_trial("tggn", 0.0, args.trials as f64); @@ -375,6 +379,7 @@ fn run_tlob_hyperopt(args: &Args) -> Result { .max_trials(args.trials) .n_initial(args.n_initial) .seed(args.seed) + .model_name("tlob") .build(); training_metrics::set_hyperopt_trial("tlob", 0.0, args.trials as f64); @@ -435,6 +440,7 @@ fn run_kan_hyperopt(args: &Args) -> Result { .max_trials(args.trials) .n_initial(args.n_initial) .seed(args.seed) + .model_name("kan") .build(); training_metrics::set_hyperopt_trial("kan", 0.0, args.trials as f64); @@ -495,6 +501,7 @@ fn run_xlstm_hyperopt(args: &Args) -> Result { .max_trials(args.trials) .n_initial(args.n_initial) .seed(args.seed) + .model_name("xlstm") .build(); training_metrics::set_hyperopt_trial("xlstm", 0.0, args.trials as f64); @@ -555,6 +562,7 @@ fn run_diffusion_hyperopt(args: &Args) -> Result { .max_trials(args.trials) .n_initial(args.n_initial) .seed(args.seed) + .model_name("diffusion") .build(); training_metrics::set_hyperopt_trial("diffusion", 0.0, args.trials as f64); diff --git a/crates/ml/src/hyperopt/optimizer.rs b/crates/ml/src/hyperopt/optimizer.rs index 7396be66d..7cc227dd5 100644 --- a/crates/ml/src/hyperopt/optimizer.rs +++ b/crates/ml/src/hyperopt/optimizer.rs @@ -92,6 +92,8 @@ pub struct ArgminOptimizer { pub(crate) seed: Option, /// Maximum iterations per restart pub(crate) max_iters_per_restart: usize, + /// Model name for Prometheus metrics progress reporting (e.g. "dqn", "ppo") + pub(crate) model_name: Option, } impl Default for ArgminOptimizer { @@ -102,6 +104,7 @@ impl Default for ArgminOptimizer { n_particles: 20, seed: None, max_iters_per_restart: 50, + model_name: None, } } } @@ -134,6 +137,7 @@ impl ArgminOptimizer { n_particles: 20, seed: None, max_iters_per_restart: 50, + model_name: None, } } @@ -296,6 +300,13 @@ impl ArgminOptimizer { &trial_counter, ¶m_names, )?; + if let Some(ref name) = self.model_name { + common::metrics::training_metrics::set_hyperopt_trial( + name, + (i + 1) as f64, + self.max_trials as f64, + ); + } } // Create trial budget observer BEFORE creating cost function @@ -310,6 +321,8 @@ impl ArgminOptimizer { param_names: param_names.clone(), bounds: bounds.clone(), observer: observer.clone(), + model_name: self.model_name.clone(), + max_trials: self.max_trials, }; // Find best initial point to start optimization @@ -702,6 +715,15 @@ impl ArgminOptimizer { current, strategy.max_concurrent_trials); concurrency.store(strategy.max_concurrent_trials, std::sync::atomic::Ordering::Relaxed); } + + if let Some(ref name) = self.model_name { + let done = trial_counter.load(std::sync::atomic::Ordering::Relaxed); + common::metrics::training_metrics::set_hyperopt_trial( + name, + done as f64, + self.max_trials as f64, + ); + } } // Create trial budget observer BEFORE creating cost function @@ -715,6 +737,8 @@ impl ArgminOptimizer { param_names: param_names.clone(), bounds: bounds.clone(), observer: observer.clone(), + model_name: self.model_name.clone(), + max_trials: self.max_trials, }; // Find best initial point @@ -847,6 +871,7 @@ pub fn optimize_with_tpe( max_trials: usize, n_initial: usize, seed: Option, + model_name: Option<&str>, ) -> Result> where M: HyperparameterOptimizable + Send, @@ -877,8 +902,10 @@ where info!(" Max Trials: {}", max_trials); info!(" Initial LHS Samples: {}", n_initial); info!(" Parameters: {}", n_params); - info!(" Gamma (good quantile): 0.25"); - info!(" EI Candidates: 100"); + let n_candidates = 256_usize.max(8 * n_params); + let gamma = if max_trials < 50 { 0.15 } else { 0.25 }; + info!(" Gamma (good quantile): {:.2}", gamma); + info!(" EI Candidates: {} (scaled for {}D)", n_candidates, n_params); let param_names = M::Params::param_names(); for (i, name) in param_names.iter().enumerate() { @@ -891,8 +918,8 @@ where n_dims: n_params, max_trials, n_initial, - gamma: 0.25, - n_candidates: 100, + gamma, + n_candidates, seed, }; let mut tpe = TpeOptimizer::new(tpe_config); @@ -932,6 +959,14 @@ where last.objective ); } + + if let Some(name) = model_name { + common::metrics::training_metrics::set_hyperopt_trial( + name, + (trial_idx + 1) as f64, + max_trials as f64, + ); + } } // Extract results @@ -990,6 +1025,8 @@ where param_names: Vec<&'static str>, bounds: Vec<(f64, f64)>, observer: crate::hyperopt::TrialBudgetObserver, + model_name: Option, + max_trials: usize, } impl CostFunction for ObjectiveFunction @@ -1085,6 +1122,14 @@ where }); } + if let Some(ref name) = self.model_name { + common::metrics::training_metrics::set_hyperopt_trial( + name, + trial_num as f64, + self.max_trials as f64, + ); + } + Ok(objective) } } @@ -1110,6 +1155,8 @@ where param_names: Vec<&'static str>, bounds: Vec<(f64, f64)>, observer: crate::hyperopt::TrialBudgetObserver, + model_name: Option, + max_trials: usize, } impl CostFunction for ParallelObjectiveFunction @@ -1200,6 +1247,14 @@ where }); } + if let Some(ref name) = self.model_name { + common::metrics::training_metrics::set_hyperopt_trial( + name, + trial_num as f64, + self.max_trials as f64, + ); + } + Ok(objective) } } @@ -1227,6 +1282,7 @@ pub struct ArgminOptimizerBuilder { n_particles: usize, seed: Option, max_iters_per_restart: usize, + model_name: Option, } impl Default for ArgminOptimizerBuilder { @@ -1244,6 +1300,7 @@ impl ArgminOptimizerBuilder { n_particles: 20, seed: None, max_iters_per_restart: 50, + model_name: None, } } @@ -1277,6 +1334,12 @@ impl ArgminOptimizerBuilder { self } + /// Set model name for per-trial Prometheus progress metrics + pub fn model_name(mut self, name: impl Into) -> Self { + self.model_name = Some(name.into()); + self + } + /// Build the optimizer /// /// # Panics @@ -1296,6 +1359,7 @@ impl ArgminOptimizerBuilder { n_particles: self.n_particles, seed: self.seed, max_iters_per_restart: self.max_iters_per_restart, + model_name: self.model_name, } } } diff --git a/crates/ml/src/hyperopt/tpe.rs b/crates/ml/src/hyperopt/tpe.rs index 46decda2a..06803539c 100644 --- a/crates/ml/src/hyperopt/tpe.rs +++ b/crates/ml/src/hyperopt/tpe.rs @@ -128,6 +128,42 @@ impl TpeOptimizer { } } + // Best-trial injection: always evaluate EI at (and around) the best known point. + // Prevents the optimizer from "forgetting" a good trial when KDE is too wide. + let best_params_opt: Option> = self + .trials + .iter() + .min_by(|a, b| { + a.objective + .partial_cmp(&b.objective) + .unwrap_or(std::cmp::Ordering::Equal) + }) + .map(|t| t.params.clone()); + + if let Some(best_params) = best_params_opt { + let ei = Self::expected_improvement_static(&best_params, &good_refs, &bad_refs, bounds); + if ei > best_ei { + best_ei = ei; + best_candidate = best_params.clone(); + } + // Small perturbations (±5% per dimension) around best trial + for _ in 0..5 { + let mut perturbed = best_params.clone(); + for (d, val) in perturbed.iter_mut().enumerate() { + if let Some(&(lo, hi)) = bounds.get(d) { + let noise = self.rng.gen_range(-0.05..0.05) * (hi - lo); + *val = (*val + noise).clamp(lo, hi); + } + } + let ei = + Self::expected_improvement_static(&perturbed, &good_refs, &bad_refs, bounds); + if ei > best_ei { + best_ei = ei; + best_candidate = perturbed; + } + } + } + best_candidate } @@ -214,8 +250,8 @@ impl TpeOptimizer { continue; } - // Compute bandwidth via Silverman's rule: h = 1.06 * std * n^(-1/5) - let bandwidth = Self::silverman_bandwidth(&values, bounds.get(d).map(|&(lo, hi)| hi - lo)); + // Scott's rule bandwidth (dimension-aware, tighter in high-D) + let bandwidth = Self::scott_bandwidth(&values, bounds.get(d).map(|&(lo, hi)| hi - lo), n_dims); // Compute log-PDF at point[d] using log-sum-exp let x = match point.get(d) { @@ -302,11 +338,11 @@ impl TpeOptimizer { // --- Private helpers --- - /// Compute Silverman's bandwidth for a set of 1-D samples. + /// Compute KDE bandwidth using Scott's rule (dimension-aware). /// - /// h = 1.06 * std * n^(-1/5), with a floor of range/100 to avoid - /// degenerate zero-bandwidth when all samples are identical. - fn silverman_bandwidth(values: &[f64], range: Option) -> f64 { + /// h = 0.7 * σ * n^(-1/(d+4)), tighter than Silverman's in high-D spaces. + /// Falls back to range/100 floor when all samples are identical. + fn scott_bandwidth(values: &[f64], range: Option, n_dims: usize) -> f64 { let n = values.len() as f64; if n < 1.0 { return 1.0; @@ -316,7 +352,9 @@ impl TpeOptimizer { let variance = values.iter().map(|&v| (v - mean) * (v - mean)).sum::() / n; let std_dev = variance.sqrt(); - let bandwidth = 1.06 * std_dev * n.powf(-0.2); + // Scott's rule with 0.7 prefactor for tighter exploitation + let exponent = -1.0 / (n_dims as f64 + 4.0); + let bandwidth = 0.7 * std_dev * n.powf(exponent); // Floor: avoid zero bandwidth (when all samples are the same) let floor = match range { @@ -350,7 +388,7 @@ impl TpeOptimizer { /// /// For each dimension independently: /// 1. Pick a random sample (kernel center) - /// 2. Add Gaussian noise with Silverman bandwidth + /// 2. Add Gaussian noise with Scott bandwidth /// 3. Clamp to bounds fn sample_from_kde(&mut self, samples: &[&Vec], bounds: &[(f64, f64)]) -> Vec { let n_dims = bounds.len(); @@ -380,7 +418,7 @@ impl TpeOptimizer { continue; } - let bandwidth = Self::silverman_bandwidth(&values, Some(hi - lo)); + let bandwidth = Self::scott_bandwidth(&values, Some(hi - lo), n_dims); // Pick a random kernel center let center_idx = self.rng.gen_range(0..values.len()); @@ -643,10 +681,10 @@ mod tests { } #[test] - fn test_tpe_silverman_bandwidth() { + fn test_tpe_scott_bandwidth() { // Identical values should produce floor bandwidth let values = vec![5.0, 5.0, 5.0]; - let bw = TpeOptimizer::silverman_bandwidth(&values, Some(10.0)); + let bw = TpeOptimizer::scott_bandwidth(&values, Some(10.0), 2); assert!( (bw - 0.1).abs() < 1e-10, "Zero-std should use floor bandwidth, got {bw}" @@ -654,7 +692,7 @@ mod tests { // Spread values should produce reasonable bandwidth let values2 = vec![0.0, 5.0, 10.0]; - let bw2 = TpeOptimizer::silverman_bandwidth(&values2, Some(10.0)); + let bw2 = TpeOptimizer::scott_bandwidth(&values2, Some(10.0), 2); assert!(bw2 > 0.1, "Spread values should have larger bandwidth: {bw2}"); assert!(bw2 < 10.0, "Bandwidth should be smaller than range: {bw2}"); }