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

gem(G12): finish predictive_coding — write backward + wire into training loop

Browse Source 
Found by V7-gem audit: predictive_coding_loss kernel existed (forward only,
per-sample MSE on consecutive trunk activations), and compute_predictive_coding_loss
Rust wrapper existed, but no backward and no caller. Pure scaffolding.

Self-supervised temporal smoothness on the enriched trunk h_s2 IS in the
"better form" by the V7 methodology — it operates at the gradient level on
the network representation, not as a reward shaping term. Worth wiring up.

Three changes:

1) New CUDA kernel `predictive_coding_backward` (experience_kernels.cu)

   Loss:    L = sum_{i=0..N-2} (lambda/SH2) * sum_j (h[i,j] - h[i+1,j])^2
   Grad:    dL/dh[i,j] = (2*lambda/SH2) * sum-of-affected-loss-terms

   Each h[i,j] appears in TWO loss terms (interior i): "current" of term i
   and "next" of term i-1. Boundaries (i=0, i=N-1) have one. Closed-form
   gradient written directly with no atomicAdd — one thread per (sample,
   feature) cell, each writes to a unique slot, plain += accumulates into
   bw_d_h_s2. Bit-deterministic.

2) Rust glue (gpu_dqn_trainer.rs)

   - Load `predictive_coding_backward` kernel via existing exp_module_for_mag
   - New field on DQNTrainer (predictive_coding_backward_kernel)
   - Extend `compute_predictive_coding_loss` to also launch backward as
     step 3 (after forward + reduce). Now the function name accurately
     describes what it does — both compute and accumulate gradient.

3) Integration (fused_training.rs::submit_aux_ops)

   Inserted the call right after `launch_recursive_confidence_backward`,
   before regime_scale_td_errors. Both spots accumulate into bw_d_h_s2 via
   plain +=, so ordering is irrelevant for correctness — what matters is
   that this runs INSIDE the aux_child CUDA-graph capture window AND
   BEFORE the trunk W_s2 → W_s1 backward GEMMs read bw_d_h_s2.

Why not also G6 (branch_independence) and G10 (temporal_consistency)?
Per V7-gem methodology — check for redundancy first:
  - G10 wants Lipschitz on Q for similar states. Spectral normalization
    (already wired on all 12 weight tensors) achieves *global* Lipschitz.
    G10 adds *local-pair* Lipschitz on top. Possibly redundant — needs
    a measurement before wiring blindly.
  - G6 wants the 4 advantage branches to stay diverse. NoisyNets already
    adds different parameter noise per layer → naturally diverse heads.
    Ensemble KL gradient pushes ensemble heads apart (different mechanism
    but similar intent). Possibly redundant for the 4-branch case.

G12 is unambiguously useful — trunk smoothness is a genuine gem with no
existing equivalent in the codebase. G6/G10 land separately if measurement
shows they add real value beyond spectral_norm + NoisyNets + ensemble KL.

Files touched:
  crates/ml/src/cuda_pipeline/experience_kernels.cu      (+45)
  crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs         (+33 / -3)
  crates/ml/src/trainers/dqn/fused_training.rs           (+9)

Verified: cargo check passes. Smoke test running to verify training is
stable with G12 active (lambda_pred=0.1 — small enough that any regression
is from a real bug, not dominance over the C51/IQN gradient).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-21 09:23:28 +02:00
parent 863eebb82a
commit e72885e8b6
3 changed files with 114 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -5065,6 +5065,70 @@ extern "C" __global__ void predictive_coding_loss(
}
}
/**
* predictive_coding_backward — gradient of the predictive_coding_loss above.
*
* Loss shape:
* L = sum_{i=0..N-2} lambda_pred * (1/SH2) * sum_j (h[i,j] - h[i+1,j])^2
*
* Gradient w.r.t. h[i,j]:
* dL/dh[i,j] = (2*lambda_pred/SH2) * sum over loss terms that involve h[i,j]
*
* Each h[i,j] appears in TWO loss terms (except boundaries):
* - as "current" in term i: contribution +diff(i,j) = (h[i,j] - h[i+1,j])
* - as "next" in term i-1: contribution -diff(i-1,j) = -(h[i-1,j] - h[i,j])
* = (h[i,j] - h[i-1,j])
*
* Combined for interior i in [1, N-2]:
* dL/dh[i,j] = (2*lambda/SH2) * ( (h[i,j] - h[i+1,j]) + (h[i,j] - h[i-1,j]) )
* = (2*lambda/SH2) * ( 2*h[i,j] - h[i-1,j] - h[i+1,j] )
*
* Boundary i=0 (only "current" term): dL/dh[0,j] = (2*lambda/SH2) * (h[0,j] - h[1,j])
* Boundary i=N-1 (only "next" term): dL/dh[N-1,j] = (2*lambda/SH2) * (h[N-1,j] - h[N-2,j])
*
* Grid: ceil(N*SH2 / 256), Block: 256. One thread per (i, j) cell.
* Output: ADDS gradient into bw_d_h_s2[i, j] (does not overwrite — caller's
* trunk backward path expects accumulation semantics).
*
* Determinism: each thread writes to a unique (i, j) slot — plain += is safe,
* no atomicAdd needed.
*/
extern "C" __global__ void predictive_coding_backward(
const float* __restrict__ h_s2, /* [N, SH2] same buffer as forward */
float* __restrict__ d_h_s2, /* [N, SH2] accumulated trunk gradient (in/out) */
int N,
int sh2,
float lambda_pred /* must match forward */
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = N * sh2;
if (idx >= total) return;
int i = idx / sh2;
int j = idx % sh2;
/* Pre-factor: 2 * lambda / SH2 */
float k = (2.0f * lambda_pred) / (float)sh2;
long long off = (long long)i * sh2 + j;
float v_self = h_s2[off];
float grad = 0.0f;
/* "current" term i: diff = (h[i] - h[i+1]) */
if (i < N - 1) {
float v_next = h_s2[(long long)(i + 1) * sh2 + j];
grad += k * (v_self - v_next);
}
/* "next" term i-1: -diff(i-1) = (h[i] - h[i-1]) */
if (i > 0) {
float v_prev = h_s2[(long long)(i - 1) * sh2 + j];
grad += k * (v_self - v_prev);
}
/* Accumulate (caller's existing trunk backward expects += semantics) */
d_h_s2[off] += grad;
}
/* ================================================================== */
/* Kernel: homeostatic_regularizer (G16) */
/* ================================================================== */

50
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -918,7 +918,10 @@ pub struct GpuDqnTrainer {
/// Scalar output buffer [1] for temporal consistency penalty.
temporal_penalty_buf: CudaSlice<f32>,
// ── G12: Predictive coding auxiliary loss ──
// ── G12: Predictive coding auxiliary loss + backward ──
/// Backward kernel: writes per-(sample, feature) gradient of the
/// predictive_coding_loss into bw_d_h_s2 via plain += (no atomicAdd).
predictive_coding_backward_kernel: CudaFunction,
/// Kernel: predictive_coding_loss — self-supervised temporal smoothness on enriched trunk.
predictive_coding_kernel: CudaFunction,
/// Per-sample loss buffer [B] for predictive coding — reduced by c51_loss_reduce.
@@ -3374,28 +3377,40 @@ impl GpuDqnTrainer {
Ok(())
}
/// G12: Predictive coding auxiliary loss — temporal smoothness on enriched trunk.
/// MSE between consecutive h_s2 samples provides a noise-free self-supervised
/// signal for the trunk. Uses save_h_s2 (enriched after mamba2_step).
/// G12: Predictive coding auxiliary loss + backward — temporal smoothness
/// on the enriched trunk. MSE between consecutive h_s2 samples gives a
/// noise-free self-supervised signal that flows back through the trunk.
///
/// Graph-safe: writes per-sample loss to predictive_per_sample_buf, then
/// reduces via c51_loss_reduce kernel. No cuMemsetD32Async or atomicAdd.
/// Three steps, all graph-safe (no memset, no atomicAdd):
/// 1. Forward: per-sample MSE → predictive_per_sample_buf [B]
/// 2. Reduce → predictive_loss_buf [1] (for monitoring readback)
/// 3. Backward: gradient on h_s2 → ACCUMULATED into bw_d_h_s2
///
/// The trunk backward (W_s2 → W_s1) reads bw_d_h_s2, so steps 1+3 effectively
/// add a regularization term to the trunk gradient. lambda_pred = 0.1 keeps
/// the predictive contribution small relative to the C51/IQN gradient.
///
/// Caller MUST invoke this AFTER the main heads have written their share
/// of bw_d_h_s2 (so we accumulate, not overwrite) and BEFORE the trunk
/// backward GEMMs (so the trunk picks up the combined gradient).
pub(crate) fn compute_predictive_coding_loss(&self, batch_size: usize) -> Result<(), MLError> {
let sh2 = self.config.shared_h2 as i32;
let save_h_s2_ptr = self.save_h_s2.raw_ptr();
let predictive_per_sample_buf_ptr = self.predictive_per_sample_buf.raw_ptr();
// Step 1: per-sample MSE loss (graph-captured kernel, no memset needed)
let lambda_pred = 0.1_f32;
// Step 1: per-sample MSE loss
unsafe {
self.stream.launch_builder(&self.predictive_coding_kernel)
.arg(&save_h_s2_ptr)
.arg(&predictive_per_sample_buf_ptr)
.arg(&(batch_size as i32))
.arg(&sh2)
.arg(&0.1_f32) // lambda_pred
.arg(&lambda_pred)
.launch(LaunchConfig::for_num_elems(batch_size as u32))
.map_err(|e| MLError::ModelError(format!("predictive_coding: {e}")))?;
}
// Step 2: deterministic reduce → predictive_loss_buf[0]
// Step 2: deterministic reduce → predictive_loss_buf[0] (monitoring)
let loss_ptr = self.predictive_loss_buf.raw_ptr();
unsafe {
self.stream
@@ -3410,6 +3425,20 @@ impl GpuDqnTrainer {
})
.map_err(|e| MLError::ModelError(format!("predictive_coding reduce: {e}")))?;
}
// Step 3: backward → ACCUMULATE gradient into bw_d_h_s2 (one thread
// per (sample, feature) cell, no atomic — unique writes).
let bw_d_h_s2_ptr = self.bw_d_h_s2.raw_ptr();
let total = (batch_size * (sh2 as usize)) as u32;
unsafe {
self.stream.launch_builder(&self.predictive_coding_backward_kernel)
.arg(&save_h_s2_ptr)
.arg(&bw_d_h_s2_ptr)
.arg(&(batch_size as i32))
.arg(&sh2)
.arg(&lambda_pred)
.launch(LaunchConfig::for_num_elems(total))
.map_err(|e| MLError::ModelError(format!("predictive_coding_backward: {e}")))?;
}
Ok(())
}
@@ -5751,6 +5780,8 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("branch_independence_penalty load: {e}")))?;
let temporal_consistency_kernel = exp_module_for_mag.load_function("temporal_consistency_penalty")
.map_err(|e| MLError::ModelError(format!("temporal_consistency_penalty load: {e}")))?;
let predictive_coding_backward_kernel = exp_module_for_mag.load_function("predictive_coding_backward")
.map_err(|e| MLError::ModelError(format!("predictive_coding_backward load: {e}")))?;
let predictive_coding_kernel = exp_module_for_mag.load_function("predictive_coding_loss")
.map_err(|e| MLError::ModelError(format!("predictive_coding_loss load: {e}")))?;
let risk_forward_kernel = exp_module_for_mag.load_function("risk_budget_forward")
@@ -7343,6 +7374,7 @@ impl GpuDqnTrainer {
temporal_consistency_kernel,
temporal_penalty_buf,
predictive_coding_kernel,
predictive_coding_backward_kernel,
predictive_per_sample_buf,
predictive_loss_buf,
q_anchor_kernel,

9
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -1457,6 +1457,15 @@ impl FusedTrainingCtx {
self.trainer.launch_recursive_confidence_backward(self.batch_size)
.map_err(|e| anyhow::anyhow!("Recursive confidence backward: {e}"))?;
// G12: Predictive coding auxiliary loss + backward.
// Self-supervised temporal smoothness on the enriched trunk h_s2 —
// adds 2*lambda*(h[i] - h[i+1])/SH2 regularization gradient into
// bw_d_h_s2. Runs HERE so the gradient lands in bw_d_h_s2 before
// the trunk backward GEMMs consume it (and before recursive
// confidence's own trunk gradient — both accumulate via plain +=).
self.trainer.compute_predictive_coding_loss(self.batch_size)
.map_err(|e| anyhow::anyhow!("Predictive coding (G12): {e}"))?;
// Ensemble diversity: heads 1..K-1 forward + KL gradient + trunk backward.
// Runs inside aux_child graph — gradients land in grad_buf before Adam.
if !self.ensemble_extra_heads.is_empty() {