gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(ml): improve TPE exploitation with Scott bandwidth, best-trial injection

Browse Source 
- Replace Silverman's bandwidth (h = 1.06σn^(-1/5)) with Scott's rule
  (h = 0.7σn^(-1/(d+4))) for tighter kernels in high-D parameter spaces
- Add best-trial injection: always evaluate EI at best known point plus
  5 small perturbations (±5%), preventing optimizer from forgetting peaks
- Scale n_candidates dynamically: max(256, 8*n_dims) instead of fixed 100
- Reduce gamma from 0.25 to 0.15 when trials < 50 for tighter exploitation
- Wire model_name through PSO/TPE paths for per-trial Prometheus metrics

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-05 23:54:30 +01:00
parent 08a927639e
commit f2938b19e8
4 changed files with 130 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
crates/ml/examples/hyperopt_baseline_rl.rs
View File

@@ -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 {

8
crates/ml/examples/hyperopt_baseline_supervised.rs
View File

@@ -136,6 +136,7 @@ fn run_tft_hyperopt(args: &Args) -> Result<Value> {
.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<Value> {
.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<Value> {
.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<Value> {
.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<Value> {
.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<Value> {
.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<Value> {
.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<Value> {
.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);

72
crates/ml/src/hyperopt/optimizer.rs
View File

@@ -92,6 +92,8 @@ pub struct ArgminOptimizer {
pub(crate) seed: Option<u64>,
/// 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<String>,
}
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,
&param_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<M>(
max_trials: usize,
n_initial: usize,
seed: Option<u64>,
model_name: Option<&str>,
) -> Result<OptimizationResult<M::Params>>
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<String>,
max_trials: usize,
}
impl<M> CostFunction for ObjectiveFunction<M>
@@ -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<String>,
max_trials: usize,
}
impl<M> CostFunction for ParallelObjectiveFunction<M>
@@ -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<u64>,
max_iters_per_restart: usize,
model_name: Option<String>,
}
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<String>) -> 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,
}
}
}

62
crates/ml/src/hyperopt/tpe.rs
View File

@@ -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<Vec<f64>> = 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>) -> 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<f64>, 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::<f64>() / 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<f64>], bounds: &[(f64, f64)]) -> Vec<f64> {
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}");
}