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

feat(ml-alpha): per-horizon BCE weighting (fixes label-correlation inflation)

Browse Source 
For seq_len K and horizon h with h ≫ K, the K position-supervised
labels in a single sequence are near-identical (sequential positions'
forward windows overlap by ~(h-1)/h). Per-position BCE therefore
treats ~K highly-correlated labels as independent samples, inflating
gradient pressure on long horizons by a factor of K.

Concretely at K=96:
  h=30   → ~3 effective samples per seq (forward windows overlap ~97%)
  h=100  → ~1                          (~99%)
  h=6000 → ~1                          (~99.98%)

Per-position supervision was paying 96× the natural signal density on
h=6000, pulling the model toward fitting noise at long horizons.

Fix: the fused BCE kernel now accepts an optional
`loss_weights[N_HORIZONS]` (nullptr → uniform = no-op). Each (k, h)
loss + grad contribution is multiplied by w_h; the normaliser is the
sum of weighted valid entries instead of the raw valid count.

`auto_horizon_weights(K, horizons)` computes `w_h = min(1.0, K/h)` so
short horizons stay at full weight and long horizons collapse to
their independent-sample density. Exposed via CLI:
  --auto-horizon-weights              # K/h auto-derived
  --horizon-weights "1,1,0.5,0.1,0.02" # explicit floats

Default behaviour is uniform (1.0) — apples-to-apples with the
in-flight qf5mj baseline. Synthetic overfit still 0.6268 → 0.1144 in
250 steps (82% drop). 77 tests pass.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-17 10:17:38 +02:00
parent 248d8fe510
commit 85ce295773
5 changed files with 120 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

79
crates/ml-alpha/cuda/bce_loss_multi_horizon.cu
View File

@@ -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;
}
}

44
crates/ml-alpha/examples/alpha_train.rs
View File

@@ -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<String>,
/// 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<f32> = spec.split(',')
.map(|s| s.trim().parse().context("horizon weight parse"))
.collect::<Result<Vec<_>>>()?;
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;

7
crates/ml-alpha/src/trainer/loss.rs
View File

@@ -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<f32> = vec![1.0; input.n_horizons];
let weights_d = upload(stream, &weights_host)?;
let mut loss_d = stream.alloc_zeros::<f32>(1).context("loss alloc")?;
let mut grad_d = stream.alloc_zeros::<f32>(total).context("grad alloc")?;
let mut valid_d = stream.alloc_zeros::<i32>(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")?; }

25
crates/ml-alpha/src/trainer/perception.rs
View File

@@ -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<CudaStream>,
@@ -140,6 +160,9 @@ pub struct PerceptionTrainer {
loss_d: CudaSlice<f32>,
/// Valid-label count for BCE normalisation.
valid_d: CudaSlice<i32>,
/// Per-horizon BCE loss weights, device-resident. Filled from
/// `cfg.horizon_weights` at construction; immutable thereafter.
loss_weights_d: CudaSlice<f32>,
/// 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::<f32>(k * N_HORIZONS)?,
labels_per_k_d: stream.alloc_zeros::<f32>(k * N_HORIZONS)?,
grad_probs_per_k_d: stream.alloc_zeros::<f32>(k * N_HORIZONS)?,
loss_weights_d: upload(&stream, &cfg.horizon_weights)?,
loss_d: stream.alloc_zeros::<f32>(1)?,
valid_d: stream.alloc_zeros::<i32>(1)?,
grad_h_carry_d: stream.alloc_zeros::<f32>(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);

1
crates/ml-alpha/tests/perception_overfit.rs
View File

@@ -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");