diff --git a/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu b/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu index f4e29997b..6be6f8958 100644 --- a/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu +++ b/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu @@ -1,85 +1,100 @@ // bce_loss_multi_horizon.cu // -// Fused multi-horizon BCE forward + backward. +// Fused multi-horizon BCE forward + backward with optional per-horizon +// loss weighting. // // Input layout (row-major): // probs [n_pos, n_horizons] — model outputs in (0, 1) // labels [n_pos, n_horizons] — binary {0.0, 1.0}; NaN = mask (drop) +// loss_weights [n_horizons] — per-horizon weight (nullptr = all 1.0) +// Down-weights horizons whose forward +// labels are correlated across K positions +// (e.g. h ≫ K → w_h ≪ 1 to avoid +// inflating gradient pressure on +// near-identical labels). // Output: -// loss_out[1] — mean BCE over valid entries -// grad_probs[n_pos, n_horizons] — d loss / d prob (scaled by 1/N_valid) +// loss_out[1] — weighted mean BCE over valid entries +// grad_probs[n_pos, n_horizons]— d loss / d prob (per-element scaled) +// valid_count_out[1] — number of non-NaN labels // -// per element: +// per element (mask m = !isnan(y), weight w = loss_weights[h]): // p = clamp(probs[i], 1e-6, 1 - 1e-6) -// y = labels[i] -// L_i = -[y log p + (1-y) log(1-p)] -// dL/dp = (p - y) / (p (1-p)) +// L_i = m * w * (-[y log p + (1-y) log(1-p)]) +// dL/dp = m * w * (p - y) / (p (1-p)) // -// total_loss = (1/N_valid) * sum_i L_i -// grad_probs[i] = (p - y) / (p (1-p) * N_valid) +// Normaliser: W_valid = sum_i m_i * w_i +// total_loss = (1/W_valid) * sum_i L_i +// grad_probs[i] = L_i'_unnormalized / W_valid // -// Block tree-reduce over a single block. n_pos * n_horizons must fit -// within the block's per-thread iteration count (this kernel is for -// minibatches up to a few thousand elements per call — well within -// reach for n_pos = 256 * seq_len = 16384 elements). For larger -// batches, the caller chunks. +// With all weights = 1.0 the kernel reduces to the original +// uniform-mean BCE (weighted sum / valid count = unweighted mean). extern "C" __global__ void bce_multi_horizon_forward_backward( - const float* __restrict__ probs, // [n_pos * n_horizons] - const float* __restrict__ labels, // [n_pos * n_horizons] + const float* __restrict__ probs, // [n_pos * n_horizons] + const float* __restrict__ labels, // [n_pos * n_horizons] + const float* __restrict__ loss_weights, // [n_horizons] (nullptr OK → 1.0) int n_pos, int n_horizons, - float* __restrict__ loss_out, // [1] - float* __restrict__ grad_probs, // [n_pos * n_horizons] - int* __restrict__ valid_count_out // [1] — number of non-NaN labels + float* __restrict__ loss_out, // [1] + float* __restrict__ grad_probs, // [n_pos * n_horizons] + int* __restrict__ valid_count_out // [1] — number of non-NaN labels ) { const int tid = threadIdx.x; const int total = n_pos * n_horizons; - __shared__ float sloss[256]; - __shared__ int svalid[256]; + __shared__ float sloss[256]; // weighted-loss numerator accumulator + __shared__ float sw_valid[256]; // weighted-mask normaliser accumulator + __shared__ int svalid[256]; // unweighted valid-label counter (reported) sloss[tid] = 0.0f; + sw_valid[tid] = 0.0f; svalid[tid] = 0; - // First pass: count valid entries and accumulate loss numerator. + // First pass: count valid entries and accumulate weighted loss + weight sum. for (int i = tid; i < total; i += blockDim.x) { const float y = labels[i]; if (isnan(y)) { grad_probs[i] = 0.0f; continue; } + const int h = i % n_horizons; + const float w = (loss_weights != nullptr) ? loss_weights[h] : 1.0f; const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f); - sloss[tid] -= y * logf(p) + (1.0f - y) * logf(1.0f - p); + sloss[tid] -= w * (y * logf(p) + (1.0f - y) * logf(1.0f - p)); + sw_valid[tid] += w; svalid[tid] += 1; } __syncthreads(); - // Block tree-reduce on sloss and svalid. + // Block tree-reduce on the three accumulators. for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (tid < s) { sloss[tid] += sloss[tid + s]; + sw_valid[tid] += sw_valid[tid + s]; svalid[tid] += svalid[tid + s]; } __syncthreads(); } if (tid == 0) { - const int n_valid = svalid[0]; - loss_out[0] = (n_valid > 0) ? sloss[0] / (float) n_valid : 0.0f; - valid_count_out[0] = n_valid; + const float w_valid = sw_valid[0]; + loss_out[0] = (w_valid > 0.0f) ? sloss[0] / w_valid : 0.0f; + valid_count_out[0] = svalid[0]; } __syncthreads(); - // Second pass: emit grad scaled by 1/n_valid. Uses the published - // valid_count_out[0] (already written by thread 0 above). - const int n_valid = valid_count_out[0]; - const float scale = (n_valid > 0) ? (1.0f / (float) n_valid) : 0.0f; + // Second pass: emit per-element grad scaled by w_h / W_valid. + // (Re-derive W_valid via shared mem instead of an extra load — we + // already have it in sw_valid[0] from the reduction above.) + const float w_valid = sw_valid[0]; + const float inv = (w_valid > 0.0f) ? (1.0f / w_valid) : 0.0f; for (int i = tid; i < total; i += blockDim.x) { const float y = labels[i]; if (isnan(y)) { continue; } + const int h = i % n_horizons; + const float w = (loss_weights != nullptr) ? loss_weights[h] : 1.0f; const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f); - grad_probs[i] = (p - y) / (p * (1.0f - p)) * scale; + grad_probs[i] = w * (p - y) / (p * (1.0f - p)) * inv; } } diff --git a/crates/ml-alpha/examples/alpha_train.rs b/crates/ml-alpha/examples/alpha_train.rs index 269358c42..75a2c72c7 100644 --- a/crates/ml-alpha/examples/alpha_train.rs +++ b/crates/ml-alpha/examples/alpha_train.rs @@ -26,7 +26,7 @@ use clap::Parser; use ml_alpha::data::loader::{MultiHorizonLoader, MultiHorizonLoaderConfig}; use ml_alpha::eval::auc::{compute_auc, AucInput}; use ml_alpha::heads::N_HORIZONS; -use ml_alpha::trainer::perception::{PerceptionTrainer, PerceptionTrainerConfig}; +use ml_alpha::trainer::perception::{auto_horizon_weights, PerceptionTrainer, PerceptionTrainerConfig}; use ml_core::device::MlDevice; use serde::Serialize; use std::path::PathBuf; @@ -84,6 +84,18 @@ struct Cli { /// val_loss improvement. 0 disables early stopping. #[arg(long, default_value_t = 3)] early_stop_patience: usize, + + /// 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)] + horizon_weights: Option, + + /// Compute per-horizon weights as `min(1, K/h)` — short horizons + /// at weight 1, long horizons down-weighted to their independent- + /// sample density inside the K-snapshot window. Recommended when + /// horizons include values larger than seq_len. + #[arg(long, default_value_t = false)] + auto_horizon_weights: bool, } #[derive(Serialize)] @@ -133,17 +145,43 @@ fn main() -> Result<()> { let dev = MlDevice::cuda(0).context("CUDA 0 init")?; tracing::info!(?dev, "MlDevice initialized"); + let horizons = [30usize, 100, 300, 1000, 6000]; + + // Resolve per-horizon BCE weights: explicit > auto > uniform. + let horizon_weights: [f32; N_HORIZONS] = if let Some(spec) = &cli.horizon_weights { + let parts: Vec = spec.split(',') + .map(|s| s.trim().parse().context("horizon weight parse")) + .collect::>>()?; + anyhow::ensure!( + parts.len() == N_HORIZONS, + "--horizon-weights expected {} floats, got {}", + N_HORIZONS, parts.len() + ); + let mut w = [0.0; N_HORIZONS]; + w.copy_from_slice(&parts); + w + } else if cli.auto_horizon_weights { + auto_horizon_weights(cli.seq_len, &horizons) + } else { + [1.0; N_HORIZONS] + }; + tracing::info!( + w_h30 = horizon_weights[0], w_h100 = horizon_weights[1], + w_h300 = horizon_weights[2], w_h1000 = horizon_weights[3], + w_h6000 = horizon_weights[4], + "per-horizon BCE weights" + ); + let trainer_cfg = PerceptionTrainerConfig { seq_len: cli.seq_len, mamba2_state_dim: cli.mamba2_state_dim, lr_cfc: cli.lr_cfc, lr_mamba2: cli.lr_mamba2, seed: cli.seed, + horizon_weights, }; let mut trainer = PerceptionTrainer::new(&dev, &trainer_cfg).context("trainer init")?; - let horizons = [30usize, 100, 300, 1000, 6000]; - let mut train_loss_running = 0.0_f32; let mut train_steps = 0usize; let mut final_val_loss = 0.0_f32; diff --git a/crates/ml-alpha/src/trainer/loss.rs b/crates/ml-alpha/src/trainer/loss.rs index c7767386f..50826dd34 100644 --- a/crates/ml-alpha/src/trainer/loss.rs +++ b/crates/ml-alpha/src/trainer/loss.rs @@ -35,6 +35,11 @@ pub fn bce_multi_horizon_loss_and_grad_gpu(dev: &MlDevice, input: &BceInput) -> let probs_d = upload(stream, &input.probs)?; let labels_d = upload(stream, &input.labels)?; + // Standalone helper invokes the kernel with uniform (all-ones) + // per-horizon weights — matches the production trainer's + // `horizon_weights: [1.0; 5]` default semantics. + let weights_host: Vec = vec![1.0; input.n_horizons]; + let weights_d = upload(stream, &weights_host)?; let mut loss_d = stream.alloc_zeros::(1).context("loss alloc")?; let mut grad_d = stream.alloc_zeros::(total).context("grad alloc")?; let mut valid_d = stream.alloc_zeros::(1).context("valid alloc")?; @@ -48,7 +53,7 @@ pub fn bce_multi_horizon_loss_and_grad_gpu(dev: &MlDevice, input: &BceInput) -> }; let mut launch = stream.launch_builder(&func); launch - .arg(&probs_d).arg(&labels_d) + .arg(&probs_d).arg(&labels_d).arg(&weights_d) .arg(&n_pos_i).arg(&n_horizons_i) .arg(&mut loss_d).arg(&mut grad_d).arg(&mut valid_d); unsafe { launch.launch(cfg).context("bce launch")?; } diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index be782721b..3f6761c4c 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -64,6 +64,12 @@ pub struct PerceptionTrainerConfig { pub lr_cfc: f32, pub lr_mamba2: f32, pub seed: u64, + /// Per-horizon BCE loss weights — multiplied into the loss + grad + /// contribution at each (position, horizon) pair. Down-weights + /// horizons whose forward labels overlap across K positions (h ≫ K + /// → near-identical labels would otherwise inflate gradient + /// pressure). All-ones gives the original uniform-mean BCE. + pub horizon_weights: [f32; N_HORIZONS], } impl Default for PerceptionTrainerConfig { @@ -74,10 +80,24 @@ impl Default for PerceptionTrainerConfig { lr_cfc: 3e-3, lr_mamba2: 1e-3, seed: 0x4242, + horizon_weights: [1.0; N_HORIZONS], } } } +/// Compute `w_h = min(1.0, K / max(h, K))` for each horizon — short +/// horizons keep full weight, long horizons get down-weighted by the +/// fraction of independent samples they actually represent in a +/// K-snapshot sequence. Used when the CLI passes `--auto-horizon-weights`. +pub fn auto_horizon_weights(seq_len: usize, horizons: &[usize; N_HORIZONS]) -> [f32; N_HORIZONS] { + let k = seq_len as f32; + let mut out = [1.0f32; N_HORIZONS]; + for (i, &h) in horizons.iter().enumerate() { + out[i] = (k / (h as f32).max(k)).min(1.0); + } + out +} + pub struct PerceptionTrainer { cfg: PerceptionTrainerConfig, stream: Arc, @@ -140,6 +160,9 @@ pub struct PerceptionTrainer { loss_d: CudaSlice, /// Valid-label count for BCE normalisation. valid_d: CudaSlice, + /// Per-horizon BCE loss weights, device-resident. Filled from + /// `cfg.horizon_weights` at construction; immutable thereafter. + loss_weights_d: CudaSlice, /// Carry of grad_h_old across the reverse-order K backward loop. /// Initialised to zero at step start; cfc_step_bwd writes its own /// grad_h_old here, which heads_bwd of the previous position adds @@ -264,6 +287,7 @@ impl PerceptionTrainer { probs_per_k_d: stream.alloc_zeros::(k * N_HORIZONS)?, labels_per_k_d: stream.alloc_zeros::(k * N_HORIZONS)?, grad_probs_per_k_d: stream.alloc_zeros::(k * N_HORIZONS)?, + loss_weights_d: upload(&stream, &cfg.horizon_weights)?, loss_d: stream.alloc_zeros::(1)?, valid_d: stream.alloc_zeros::(1)?, grad_h_carry_d: stream.alloc_zeros::(n_hid)?, @@ -547,6 +571,7 @@ impl PerceptionTrainer { let mut launch = self.stream.launch_builder(&self.bce_fn); launch .arg(&self.probs_per_k_d).arg(&self.labels_per_k_d) + .arg(&self.loss_weights_d) .arg(&n_pos_i).arg(&n_h_i) .arg(&mut self.loss_d).arg(&mut self.grad_probs_per_k_d) .arg(&mut self.valid_d); diff --git a/crates/ml-alpha/tests/perception_overfit.rs b/crates/ml-alpha/tests/perception_overfit.rs index 68d5226fd..26ce44124 100644 --- a/crates/ml-alpha/tests/perception_overfit.rs +++ b/crates/ml-alpha/tests/perception_overfit.rs @@ -70,6 +70,7 @@ fn stacked_trainer_loss_shrinks_on_constant_signal() { lr_cfc: 3e-3, lr_mamba2: 1e-3, seed: 0x4242, + horizon_weights: [1.0; 5], }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");