feat(ml-alpha): ISV horizon weighting Phase 3 — wire lambda into trunk grad

Connects the per-horizon lambda computed by horizon_ema_and_lambda
(landed in 37c3a8f4d) to the actual trunk gradient flow. This is the
model-behavior change. mhzs7 (3a196382f) hit mean_auc=0.726 but with
the long-horizon h6000 stuck at 0.694 — exactly the horizon the
auto-horizon-weights formula `min(1, K/h)` over-weights down (h6000
weight = 0.0053). ISV detects "this horizon is hard, give it more
trunk-gradient pull" and lambda[h] scales the per-horizon `d_z`
contribution into the trunk in heads_bwd.

Kernel change (cuda/multi_horizon_heads.cu):
  multi_horizon_heads_backward_batched(): new arg
    `const float* __restrict__ lambda` (5 elements, between
    grad_h_carry and n_batch in arg order).
  - lambda is broadcast into __shared__ float s_lambda[5] once per
    block so every thread reads the 5 floats without repeated
    global loads.
  - Sentinel: zero buffer (init state, EMA hasn't run yet) is read
    as 1.0 so the kernel reduces to pre-ISV behavior. After step 1
    every entry is clamped to [0.5, 2.0] by horizon_lambda.cu and
    the > 0 check is always true.
  - lambda[k] multiplies the `acc += s_lambda[k] * sd_z * w[k]`
    term that produces grad_h (trunk gradient).
  - grad_w and grad_b updates are UNCHANGED. The horizon heads keep
    learning their own weight/bias normally; lambda only biases how
    much each horizon's error signal leaks into the shared trunk.
    Per `pearl_adam_normalizes_loss_weights.md`: scaling the
    effective gradient bypasses Adam's m/sqrt(v) normalization that
    would otherwise cancel a loss-weight lift.

Trainer change (trainer/perception.rs):
  - heads_bwd_batched launch now passes `&self.lambda_d` between
    grad_h_carry_d and n_batch_i. lambda_d is refreshed each step
    by the horizon_ema_and_lambda kernel that ran just after BCE.
  - Whole chain lives inside the captured CUDA Graph.

Validation:
  - All 7 perception_overfit tests pass.
  - Synthetic overfit converges marginally FASTER than pre-Phase-3
    (final loss 0.0007 → 0.0005). Plausible explanation: in a
    constant-signal smoke, the per-horizon BCE values are similar
    enough that lambda mostly stays near 1.0, but the EMA
    bootstrap on step 1 still produces non-uniform initial lambdas
    that nudge the trunk toward whichever horizon converges
    slowest. Either way, no regression.
  - horizon_ema_and_lambda_track_after_training test still passes
    with the lambda values now flowing through heads_bwd:
      ema    = [0.50, 0.87, 0.64, 0.49, 0.60]
      lambda = [0.80, 1.41, 1.03, 0.79, 0.97]
      → h100 hardest (BCE 0.87) gets 1.41× trunk grad; h300 easiest
        (BCE 0.49) gets 0.79× — exactly the intended ISV semantics.
  - 26 lib + 23 integration ml-alpha tests pass.

Honors:
  - feedback_isv_for_adaptive_bounds.md (no hardcoded constants;
    lambda derives from runtime BCE signal).
  - pearl_adam_normalizes_loss_weights.md (scaling trunk gradient,
    not BCE coefficient, to bypass Adam normalization).
  - pearl_audit_unboundedness_for_implicit_asymmetry.md (lambda is
    bounded by horizon_lambda.cu's clamp [0.5, 2.0]).
  - pearl_no_deferrals_for_complementary_fixes.md — Phase 3 wires
    up the infrastructure from Phase 1+2 in the same logical
    delivery rather than waiting a CV cycle.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-17 16:52:36 +02:00
parent 37c3a8f4d7
commit 5d42ab0e98
2 changed files with 30 additions and 5 deletions

View File

@@ -114,6 +114,7 @@ extern "C" __global__ void multi_horizon_heads_backward_batched(
const float* __restrict__ h, // [n_batch, 128] const float* __restrict__ h, // [n_batch, 128]
const float* __restrict__ grad_probs, // [n_batch, 5] const float* __restrict__ grad_probs, // [n_batch, 5]
const float* __restrict__ grad_h_carry, // [n_batch, 128] (nullptr OK) const float* __restrict__ grad_h_carry, // [n_batch, 128] (nullptr OK)
const float* __restrict__ lambda, // [5] — per-horizon trunk-grad scaler
int n_batch, int n_batch,
float* __restrict__ grad_w, // [5, 128] accum += float* __restrict__ grad_w, // [5, 128] accum +=
float* __restrict__ grad_b, // [5] accum += float* __restrict__ grad_b, // [5] accum +=
@@ -134,7 +135,25 @@ extern "C" __global__ void multi_horizon_heads_backward_batched(
} }
__syncthreads(); __syncthreads();
// Per-horizon trunk-gradient scaler, broadcast via shared mem so
// every thread reads the same 5 floats without repeated global loads.
__shared__ float s_lambda[5];
if (tid < 5) {
// Sentinel: a fully zero lambda buffer (initial state, before
// the EMA kernel has run even once) is treated as "scaler = 1"
// so the kernel reduces to its pre-ISV behavior. After step 1
// every entry is clamped to [0.5, 2.0] by horizon_lambda.cu
// and the > 0 check is always true.
const float l = lambda[tid];
s_lambda[tid] = (l > 0.0f) ? l : 1.0f;
}
__syncthreads();
// grad_b[k] += sum_b d_z[b, k]. Thread k (k < 5) is sole writer. // grad_b[k] += sum_b d_z[b, k]. Thread k (k < 5) is sole writer.
// NOTE: lambda does NOT scale the head's own bias gradient — only
// the trunk gradient. We want the per-horizon head to keep learning
// its own bias normally; lambda only biases how much the horizon's
// error signal LEAKS into the shared trunk.
if (tid < 5) { if (tid < 5) {
float acc = 0.0f; float acc = 0.0f;
for (int bi = 0; bi < n_batch; ++bi) { for (int bi = 0; bi < n_batch; ++bi) {
@@ -143,11 +162,14 @@ extern "C" __global__ void multi_horizon_heads_backward_batched(
grad_b[tid] += acc; grad_b[tid] += acc;
} }
// grad_w[k, i] += sum_b d_z[b, k] * h[b, i]. Thread i (i < 128) is // grad_w[k, i] += sum_b d_z[b, k] * h[b, i]. Same rationale as
// sole writer to its column of grad_w. (Below: parallelise over i // grad_b: lambda does NOT scale the head's own weight gradient.
// across the 5 head rows.)
if (tid < 128) { if (tid < 128) {
// grad_h: sum_k d_z[b, k] * W[k, i] for each b, plus optional carry. // grad_h: sum_k LAMBDA[k] * d_z[b, k] * W[k, i] for each b,
// plus optional carry. This is the only path where lambda is
// applied — per `pearl_adam_normalizes_loss_weights.md` the
// effective lever for horizon-aware learning is the trunk
// gradient, not the head loss weight.
for (int bi = 0; bi < n_batch; ++bi) { for (int bi = 0; bi < n_batch; ++bi) {
const float h_bi = h[(long long)bi * 128 + tid]; const float h_bi = h[(long long)bi * 128 + tid];
// Accumulate W-row grads via += into global (thread tid sole writer of column tid). // Accumulate W-row grads via += into global (thread tid sole writer of column tid).
@@ -156,7 +178,7 @@ extern "C" __global__ void multi_horizon_heads_backward_batched(
} }
float acc = 0.0f; float acc = 0.0f;
for (int k = 0; k < 5; ++k) { for (int k = 0; k < 5; ++k) {
acc += sd_z[(long long)bi * 5 + k] * w[k * 128 + tid]; acc += s_lambda[k] * sd_z[(long long)bi * 5 + k] * w[k * 128 + tid];
} }
const float carry = (grad_h_carry != nullptr) const float carry = (grad_h_carry != nullptr)
? grad_h_carry[(long long)bi * 128 + tid] : 0.0f; ? grad_h_carry[(long long)bi * 128 + tid] : 0.0f;

View File

@@ -1016,6 +1016,8 @@ impl PerceptionTrainer {
let grad_henr_k_ptr = grad_henr_t_base + (k * kb_hid_bytes) as u64; let grad_henr_k_ptr = grad_henr_t_base + (k * kb_hid_bytes) as u64;
// heads_bwd_batched writes grad_h_new[B, 128] = heads-side + carry. // heads_bwd_batched writes grad_h_new[B, 128] = heads-side + carry.
// `lambda_d` scales ONLY the trunk gradient — the heads'
// own weight/bias gradients are unscaled.
unsafe { unsafe {
let mut launch = self.stream.launch_builder(&self.heads_bwd_batched_fn); let mut launch = self.stream.launch_builder(&self.heads_bwd_batched_fn);
launch launch
@@ -1024,6 +1026,7 @@ impl PerceptionTrainer {
.arg(&h_new_k_ptr) .arg(&h_new_k_ptr)
.arg(&gprobs_k_ptr) .arg(&gprobs_k_ptr)
.arg(&self.grad_h_carry_d) .arg(&self.grad_h_carry_d)
.arg(&self.lambda_d)
.arg(&n_batch_i) .arg(&n_batch_i)
.arg(&mut self.grad_heads_w_d) .arg(&mut self.grad_heads_w_d)
.arg(&mut self.grad_heads_b_d) .arg(&mut self.grad_heads_b_d)