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: clamped to `[LAMBDA_FLOOR, LAMBDA_CEILING]` per
// `pearl_audit_unboundedness_for_implicit_asymmetry.md` — the ratio
// loss_ema_h / mean_loss_ema is unbounded above when one horizon
// stalls. Capping at 2× prevents winner-take-all amplification per
// `pearl_controller_amplifies_dominant_magnitude_trap.md`.

#define N_HORIZONS_LAMBDA 5
#define ALPHA_FIXED       0.1f
#define LAMBDA_FLOOR      0.5f
#define LAMBDA_CEILING    2.0f

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;
    // mean > 1e-12 is essentially always true once the first step
    // has run (BCE for a random init is ~0.69), but guard anyway.
    const float inv_mean = (mean > 1e-12f) ? (1.0f / mean) : 0.0f;
    #pragma unroll
    for (int h = 0; h < N_HORIZONS_LAMBDA; ++h) {
        const float ratio = (inv_mean > 0.0f) ? new_ema[h] * inv_mean : 1.0f;
        lambda[h] = fmaxf(LAMBDA_FLOOR, fminf(LAMBDA_CEILING, ratio));
    }
}