foxhunt/crates/ml-alpha/cuda/horizon_lambda.cu

// horizon_lambda.cu — ISV-driven per-horizon gradient scaler.
//
// Step 1: maintain an EMA of the UNWEIGHTED per-horizon BCE loss that
//         the BCE kernel emits each training step into
//         `loss_per_horizon[N_HORIZONS]`.
// Step 2: convert the EMA into a per-horizon multiplicative lambda
//         used by the backward path to scale how strongly each
//         horizon influences the shared trunk gradient.
//
// Why ISV: the current static `auto-horizon-weights` formula
// (`min(1, K/h)`) is a closed-form heuristic that ignores actual
// per-horizon learning difficulty. Empirically mhzs7 spent most of
// training over-weighting short horizons (whose label correlation
// within the K-snapshot window dominates the gradient signal) while
// h6000 — the deployment-relevant multi-minute horizon — stayed at
// AUC≈0.69. Tracking per-horizon BCE directly lets the lambda boost
// the horizons that the model is currently failing to learn, without
// hand-tuned constants.
//
// Why not just lift BCE coefficients: per
// `pearl_adam_normalizes_loss_weights.md`, Adam's m/sqrt(v) cancels
// per-loss weight lifts (SP13: 13× aux_w produced only 0.6%/epoch
// divergence). The effective lever is to scale the GRADIENT into the
// shared trunk, not the loss aggregate. heads_bwd will multiply the
// per-horizon `d_z` contribution by lambda[h] before accumulating
// into `grad_h`, bypassing Adam normalization.
//
// First-observation bootstrap: loss_ema is zero-initialised; the
// kernel detects `prev <= 0` and replaces (rather than blends) on
// the first step. After that it uses a fixed α — Wiener-optimal α is
// a Phase 3 follow-up; for now a conservative 0.1 keeps the EMA
// stable across training noise.
//
// Lambda safety: ASYMMETRIC clamp `[1.0, LAMBDA_CEILING]` per
// `pearl_audit_unboundedness_for_implicit_asymmetry.md` — boost-only.
// Lambda derivation: Z-SCORE NORMALISED, not raw ratio. The ratio
// formula `ema_h / mean(ema)` produced lambdas in 0.97-1.03 in our
// data because per-horizon BCE clusters tightly (range 0.04 abs)
// while the mean is ~0.65 — so the controller barely engaged
// (asymmetric ceiling 2.0 was never approached, max observed
// lambda ~1.04). Z-score `z_h = (ema_h - mean) / std(ema)` makes
// the spread scale-invariant per
// `pearl_zscore_normalization_for_magnitude_asymmetric_signals.md`,
// then `lambda = 1.0 + Z_SCALE * z` fills the clamp envelope: a
// horizon 1σ above mean gets lambda 1.5 (with Z_SCALE=0.5), 2σ
// saturates at the ceiling. Boost-only floor at 1.0 still applies.

#define N_HORIZONS_LAMBDA 5
#define ALPHA_FIXED       0.1f
#define LAMBDA_FLOOR      1.0f    // boost-only: never demote a horizon below uniform
#define LAMBDA_CEILING    2.0f
#define Z_SCALE           0.5f    // 1σ above mean → lambda 1.5; 2σ → ceiling

extern "C" __global__ void horizon_ema_and_lambda(
    const float* __restrict__ loss_per_horizon,  // [5] — current step UNWEIGHTED BCE
    float* __restrict__ loss_ema,                 // [5] — EMA state (read + write)
    float* __restrict__ lambda                    // [5] — output multiplier
) {
    if (threadIdx.x != 0 || blockIdx.x != 0) return;

    float new_ema[N_HORIZONS_LAMBDA];
    float sum = 0.0f;
    #pragma unroll
    for (int h = 0; h < N_HORIZONS_LAMBDA; ++h) {
        const float cur = loss_per_horizon[h];
        const float prev = loss_ema[h];
        // Sentinel = 0 ⇒ first observation replaces directly.
        // `prev <= 0` is safer than `== 0` under --use_fast_math.
        new_ema[h] = (prev <= 0.0f) ? cur : (prev + ALPHA_FIXED * (cur - prev));
        loss_ema[h] = new_ema[h];
        sum += new_ema[h];
    }

    const float mean = sum / (float)N_HORIZONS_LAMBDA;
    // Per-horizon std (population, not sample — N=5 is small + fixed).
    // EPSILON guards the first-step case where all EMAs are equal
    // (z would be 0/0); under that condition lambda falls back to 1.0
    // via the asymmetric floor.
    float ssq = 0.0f;
    #pragma unroll
    for (int h = 0; h < N_HORIZONS_LAMBDA; ++h) {
        const float diff = new_ema[h] - mean;
        ssq += diff * diff;
    }
    const float std = sqrtf(ssq / (float)N_HORIZONS_LAMBDA + 1e-12f);
    const float inv_std = (std > 1e-6f) ? (1.0f / std) : 0.0f;
    #pragma unroll
    for (int h = 0; h < N_HORIZONS_LAMBDA; ++h) {
        const float z = (new_ema[h] - mean) * inv_std;
        const float raw = 1.0f + Z_SCALE * z;
        lambda[h] = fmaxf(LAMBDA_FLOOR, fminf(LAMBDA_CEILING, raw));
    }
}