From 8736afb16caa7302b98df51cbfb586a66f3964d3 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 8 Apr 2026 18:37:55 +0200 Subject: [PATCH] =?UTF-8?q?feat(magnitude):=20advantage=20standardization?= =?UTF-8?q?=20+=205=C3=97=20entropy=20boost=20=E2=80=94=20prevents=20Q-val?= =?UTF-8?q?ue=20and=20distribution=20collapse?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/c51_grad_kernel.cu | 7 ++-- .../ml/src/cuda_pipeline/c51_loss_kernel.cu | 32 ++++++++++++++----- .../ml/src/cuda_pipeline/mse_grad_kernel.cu | 7 ++++ .../ml/src/cuda_pipeline/mse_loss_kernel.cu | 32 ++++++++++++++----- 4 files changed, 60 insertions(+), 18 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu index e6dfe0489..201918af9 100644 --- a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu @@ -41,10 +41,13 @@ extern "C" __global__ void c51_grad_kernel( * Float exp() handles full range — no bf16 overflow */ float d_combined = isw * (expf(lp) - proj); - /* Entropy regularization: d/d_logits(-coeff * H(p)) = coeff * (1 + lp) */ + /* Entropy regularization — magnitude branch (d==1) gets 5× boost to prevent + * atom distribution collapse. Standardized advantages prevent Q-value scale + * collapse; this prevents atom sharpness collapse. Two-pronged defense. */ if (entropy_coeff > 0.0f) { + float ent_scale = (d == 1) ? 5.0f : 1.0f; float lp_clamped = fmaxf(lp, -10.0f); - d_combined += entropy_coeff * (1.0f + lp_clamped); + d_combined += ent_scale * entropy_coeff * (1.0f + lp_clamped); } /* d_value_logits is f32 — native atomicAdd, no overflow risk. */ diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index 0e3f429fd..b79f0ef1c 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -335,14 +335,26 @@ extern "C" __global__ void c51_loss_batched( shmem_adv[i] = on_adv_row[d][i]; /* already f32 */ __syncthreads(); - /* Dueling: Q[j] = V[j] + A[a_d,j] - mean_a(A[*,j]) - * All 4 branches now have small action counts — mean subtraction is safe. */ + /* Dueling: Q[j] = V[j] + A'[a_d,j] where A' is centered (all branches) + * and standardized to unit variance (magnitude branch d==1 only). + * Standardization prevents winner-take-all Q-value collapse where Small + * dominates due to lower PnL variance → tighter C51 distributions. */ for (int j = tid; j < num_atoms; j += BLOCK_THREADS) { float a_mean = 0.0f; for (int a = 0; a < n_d; a++) a_mean += shmem_adv[a * num_atoms + j]; a_mean /= (float)n_d; - shmem_lp[j] = shmem_val[j] + shmem_adv[a_d * num_atoms + j] - a_mean; + float centered = shmem_adv[a_d * num_atoms + j] - a_mean; + if (d == 1) { + float sq_sum = 0.0f; + for (int a = 0; a < n_d; a++) { + float diff = shmem_adv[a * num_atoms + j] - a_mean; + sq_sum += diff * diff; + } + float a_std = sqrtf(sq_sum / (float)n_d + 1e-12f); + centered /= (a_std + 1e-6f); + } + shmem_lp[j] = shmem_val[j] + centered; } __syncthreads(); @@ -352,17 +364,21 @@ extern "C" __global__ void c51_loss_batched( * non-Flat direction, weighted at 0.3×. This gives the magnitude branch * gradient signal on ALL samples, not just the 42% with directional actions. */ if (d == 1 && branch_action[0] == 1) { - /* Magnitude advantage is already in shmem_adv. Compute counterfactual - * Q = V + 0.3 * (A_mag[a_d] - mean(A_mag)) instead of Q = V only. - * The 0.3× weight prevents counterfactual from dominating real gradient. */ float cf_weight = 0.3f; for (int j = tid; j < num_atoms; j += BLOCK_THREADS) { float a_mean = 0.0f; for (int a = 0; a < n_d; a++) a_mean += shmem_adv[a * num_atoms + j]; a_mean /= (float)n_d; - shmem_lp[j] = shmem_val[j] - + cf_weight * (shmem_adv[a_d * num_atoms + j] - a_mean); + float centered = shmem_adv[a_d * num_atoms + j] - a_mean; + float sq_sum = 0.0f; + for (int a = 0; a < n_d; a++) { + float diff = shmem_adv[a * num_atoms + j] - a_mean; + sq_sum += diff * diff; + } + float a_std = sqrtf(sq_sum / (float)n_d + 1e-12f); + centered /= (a_std + 1e-6f); + shmem_lp[j] = shmem_val[j] + cf_weight * centered; } } __syncthreads(); diff --git a/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu b/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu index 28ac45efd..8c0320cce 100644 --- a/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu @@ -45,6 +45,13 @@ extern "C" __global__ void mse_grad_kernel( /* Gradient of MSE loss through softmax expectation (float arithmetic) */ float d_combined = isw * td_error * p_j * (z_j - e_q); + /* Magnitude entropy boost: prevent atom distribution collapse for d==1. + * MSE path has no entropy param — hardcoded 0.005 = 5× base (0.001). */ + if (d == 1) { + float lp_approx = fmaxf(logf(fmaxf(p_j, 1e-8f)), -10.0f); + d_combined += 0.005f * (1.0f + lp_approx); + } + /* Route through dueling: d_value[b,j] += d_combined. * d_value_logits is f32 — native atomicAdd, no overflow risk. */ atomicAdd(&d_value_logits[b * num_atoms + j], d_combined); diff --git a/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu b/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu index 4b5884271..3f3d4bbd2 100644 --- a/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu @@ -283,14 +283,26 @@ extern "C" __global__ void mse_loss_batched( shmem_adv[i] = on_adv_row[d][i]; __syncthreads(); - /* Dueling: Q[j] = V[j] + A[a_d,j] - mean_a(A[*,j]) - * All 4 branches now have small action counts — mean subtraction is safe. */ + /* Dueling: Q[j] = V[j] + A'[a_d,j] where A' is centered (all branches) + * and standardized to unit variance (magnitude branch d==1 only). + * Standardization prevents winner-take-all Q-value collapse where Small + * dominates due to lower PnL variance → tighter C51 distributions. */ for (int j = tid; j < num_atoms; j += BLOCK_THREADS) { float a_mean = 0.0f; for (int a = 0; a < n_d; a++) a_mean += shmem_adv[a * num_atoms + j]; a_mean /= (float)n_d; - shmem_lp[j] = shmem_val[j] + shmem_adv[a_d * num_atoms + j] - a_mean; + float centered = shmem_adv[a_d * num_atoms + j] - a_mean; + if (d == 1) { + float sq_sum = 0.0f; + for (int a = 0; a < n_d; a++) { + float diff = shmem_adv[a * num_atoms + j] - a_mean; + sq_sum += diff * diff; + } + float a_std = sqrtf(sq_sum / (float)n_d + 1e-12f); + centered /= (a_std + 1e-6f); + } + shmem_lp[j] = shmem_val[j] + centered; } __syncthreads(); @@ -300,17 +312,21 @@ extern "C" __global__ void mse_loss_batched( * non-Flat direction, weighted at 0.3×. This gives the magnitude branch * gradient signal on ALL samples, not just the 42% with directional actions. */ if (d == 1 && branch_action[0] == 1) { - /* Magnitude advantage is already in shmem_adv. Compute counterfactual - * Q = V + 0.3 * (A_mag[a_d] - mean(A_mag)) instead of Q = V only. - * The 0.3× weight prevents counterfactual from dominating real gradient. */ float cf_weight = 0.3f; for (int j = tid; j < num_atoms; j += BLOCK_THREADS) { float a_mean = 0.0f; for (int a = 0; a < n_d; a++) a_mean += shmem_adv[a * num_atoms + j]; a_mean /= (float)n_d; - shmem_lp[j] = shmem_val[j] - + cf_weight * (shmem_adv[a_d * num_atoms + j] - a_mean); + float centered = shmem_adv[a_d * num_atoms + j] - a_mean; + float sq_sum = 0.0f; + for (int a = 0; a < n_d; a++) { + float diff = shmem_adv[a * num_atoms + j] - a_mean; + sq_sum += diff * diff; + } + float a_std = sqrtf(sq_sum / (float)n_d + 1e-12f); + centered /= (a_std + 1e-6f); + shmem_lp[j] = shmem_val[j] + cf_weight * centered; } } __syncthreads();