gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: Layer 1+2 — C51 loss clamp + label smoothing (prevents 242K explosion)

Browse Source 
Layer 1: Per-sample CE clamped to MAX_PER_SAMPLE_CE=50 before IS weight.
Breaks PER feedback loop: high loss → high priority → high IS weight → repeat.

Layer 2: Bellman target smoothed with ε=0.01 uniform mix after projection.
Prevents any atom from having zero probability → no log(near-zero) in CE.

Root cause: C51 cross-entropy is unbounded when target and predicted
distributions are maximally misaligned. PER amplifies pathological samples.
These two layers cap the maximum possible CE and prevent the extreme
misalignment from occurring.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-23 20:45:11 +01:00
parent 909650f2a1
commit a51345a141
1 changed files with 27 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
View File

@@ -22,6 +22,12 @@
*/
/* ── Compile-time defaults (overridden by NVRTC injection) ───────────── */
#ifndef MAX_PER_SAMPLE_CE
#define MAX_PER_SAMPLE_CE 50.0f /* Layer 1: 10x converged loss — breaks PER feedback loop */
#endif
#ifndef LABEL_SMOOTHING_EPS
#define LABEL_SMOOTHING_EPS 0.01f /* Layer 2: 1% uniform mix — prevents target collapse */
#endif
#ifndef NUM_ATOMS
#define NUM_ATOMS 51
#endif
@@ -567,7 +573,19 @@ extern "C" __global__ void c51_loss_batched(
block_bellman_project(shmem_proj, reward, done, gamma,
shmem_lp, shmem_support, shmem_reduce, tid);
/* shmem_lp[0..NUM_ATOMS] = projected target distribution */
__syncthreads();
/* Layer 2: Label smoothing — mix ε uniform into projected target.
* Prevents any atom from having zero probability → no log(0) in CE. */
{
const float ls_uniform = LABEL_SMOOTHING_EPS / (float)NUM_ATOMS;
const float ls_scale = 1.0f - LABEL_SMOOTHING_EPS;
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
shmem_lp[j] = ls_scale * shmem_lp[j] + ls_uniform;
__syncthreads();
}
/* shmem_lp[0..NUM_ATOMS] = smoothed projected target distribution */
/* Save projected for backward kernel */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
@@ -594,9 +612,15 @@ extern "C" __global__ void c51_loss_batched(
/* ── Write outputs (thread 0 only to avoid races) ────────────── */
if (tid == 0) {
float weighted_loss = avg_ce * is_weight;
/* Layer 1: clamp per-sample CE to prevent PER feedback loop.
* High loss → high PER priority → high IS weight → amplified loss → repeat.
* Clamping at 50.0 (10x converged loss of ~5) breaks the loop.
* The gradient itself is bounded (softmax - target ∈ [0,1]),
* so clamping the LOSS only affects PER priority, not gradient direction. */
float clamped_ce = fminf(avg_ce, MAX_PER_SAMPLE_CE);
float weighted_loss = clamped_ce * is_weight;
per_sample_loss[sample_id] = weighted_loss;
td_errors[sample_id] = avg_ce; /* unweighted — used by PER for priorities */
td_errors[sample_id] = clamped_ce; /* PER sees clamped too */
atomicAdd(total_loss, weighted_loss / (float)batch_size);
}
}