diff --git a/crates/ml-alpha/cuda/multi_horizon_heads.cu b/crates/ml-alpha/cuda/multi_horizon_heads.cu index 7ab21731e..83f54fe54 100644 --- a/crates/ml-alpha/cuda/multi_horizon_heads.cu +++ b/crates/ml-alpha/cuda/multi_horizon_heads.cu @@ -114,6 +114,7 @@ extern "C" __global__ void multi_horizon_heads_backward_batched( const float* __restrict__ h, // [n_batch, 128] const float* __restrict__ grad_probs, // [n_batch, 5] const float* __restrict__ grad_h_carry, // [n_batch, 128] (nullptr OK) + const float* __restrict__ lambda, // [5] — per-horizon trunk-grad scaler int n_batch, float* __restrict__ grad_w, // [5, 128] accum += float* __restrict__ grad_b, // [5] accum += @@ -134,7 +135,25 @@ extern "C" __global__ void multi_horizon_heads_backward_batched( } __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. + // 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) { float acc = 0.0f; 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_w[k, i] += sum_b d_z[b, k] * h[b, i]. Thread i (i < 128) is - // sole writer to its column of grad_w. (Below: parallelise over i - // across the 5 head rows.) + // grad_w[k, i] += sum_b d_z[b, k] * h[b, i]. Same rationale as + // grad_b: lambda does NOT scale the head's own weight gradient. 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) { const float h_bi = h[(long long)bi * 128 + 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; 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) ? grad_h_carry[(long long)bi * 128 + tid] : 0.0f; diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index d8a123607..40619d2e8 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -1016,6 +1016,8 @@ impl PerceptionTrainer { 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. + // `lambda_d` scales ONLY the trunk gradient — the heads' + // own weight/bias gradients are unscaled. unsafe { let mut launch = self.stream.launch_builder(&self.heads_bwd_batched_fn); launch @@ -1024,6 +1026,7 @@ impl PerceptionTrainer { .arg(&h_new_k_ptr) .arg(&gprobs_k_ptr) .arg(&self.grad_h_carry_d) + .arg(&self.lambda_d) .arg(&n_batch_i) .arg(&mut self.grad_heads_w_d) .arg(&mut self.grad_heads_b_d)