fix(noisy): Bug #4 — sigma_mean returns effective σ, not raw tensor (Task 2.5)

HEALTH_DIAG `sigma_mag` / `sigma_dir` were reporting raw weight_sigma
mean (constant 0.0320 across schedules) because reset_noise_with_sigma
only scaled the noise epsilon samples, not the underlying σ tensor.

Fix: added current_sigma_scale: f32 field on NoisyLinear (default 1.0),
updated in reset_noise_with_sigma, multiplied through in sigma_mean()
accessor. Also propagated through copy_params_from so target-net sync
does not reset the reported effective σ.

The HEALTH_DIAG field now reflects the scheduled effective σ, enabling
H7 detection signal to actually observe schedule attenuation.

Closes Track 4 E2 TUNE finding from Phase 1 triage.
This commit is contained in:
jgrusewski
2026-04-22 15:50:02 +02:00
parent 4c3806da9c
commit ef165e8961

View File

@@ -71,6 +71,12 @@ pub struct NoisyLinear {
in_features: usize,
out_features: usize,
// Task 2.5 Bug #4: effective sigma scale applied by the scheduler via
// `reset_noise_with_sigma`. Retained so `sigma_mean()` can report the
// *effective* σ (scaled), not the raw σ tensor constant. Default 1.0
// means "no scheduler attenuation", matching reset_noise()'s behaviour.
current_sigma_scale: f32,
// CUDA stream for GPU operations
stream: Arc<CudaStream>,
}
@@ -145,6 +151,7 @@ impl NoisyLinear {
bias_epsilon,
in_features,
out_features,
current_sigma_scale: 1.0,
stream,
})
}
@@ -163,6 +170,11 @@ impl NoisyLinear {
// where f(x) = sign(x) * sqrt(|x|)
let scale = sigma_scale as f32;
// Task 2.5 Bug #4: record scheduler attenuation so HEALTH_DIAG
// `sigma_mean` reports effective σ (= raw σ × current scale),
// not the constant raw weight_sigma tensor mean.
self.current_sigma_scale = scale;
// Generate raw Gaussian noise on host, apply f(x) = sign(x) * sqrt(|x|)
let apply_f = |x: f32| -> f32 {
let sign = if x >= 0.0 { 1.0_f32 } else { -1.0_f32 };
@@ -339,9 +351,16 @@ impl NoisyLinear {
[&self.weight_sigma, &self.bias_sigma]
}
/// Task 0.6 — mean of |sigma| across weight + bias σ tensors. Used by
/// HEALTH_DIAG to track NoisyNets exploration pressure per branch
/// (H7 detection signal). Sync stream-readback; epoch-boundary only.
/// Task 0.6 — mean of |sigma| across weight + bias σ tensors, scaled by
/// the scheduler's current sigma_scale. Used by HEALTH_DIAG to track
/// NoisyNets exploration pressure per branch (H7 detection signal).
/// Sync stream-readback; epoch-boundary only.
///
/// Task 2.5 Bug #4: result is multiplied by `current_sigma_scale` so it
/// reports effective σ (raw σ × scheduler scale), not the constant raw
/// weight_sigma tensor mean. `reset_noise_with_sigma` scales the noise
/// epsilon samples; the σ tensor itself stays fixed, so without this
/// multiplier the diag reads constant 0.0320 regardless of schedule.
pub fn sigma_mean(&self) -> Result<f32, MLError> {
let n_w = self.weight_sigma.len();
let n_b = self.bias_sigma.len();
@@ -358,7 +377,8 @@ impl NoisyLinear {
.map_err(|e| MLError::ModelError(format!("sigma_mean bias dtoh: {e}")))?;
let sum_w: f64 = h_w.iter().map(|x| x.abs() as f64).sum();
let sum_b: f64 = h_b.iter().map(|x| x.abs() as f64).sum();
Ok(((sum_w + sum_b) / ((n_w + n_b) as f64)) as f32)
let raw_mean = ((sum_w + sum_b) / ((n_w + n_b) as f64)) as f32;
Ok(raw_mean * self.current_sigma_scale)
}
/// Disable noise for evaluation (use mean parameters only)
@@ -387,6 +407,9 @@ impl NoisyLinear {
Self::dtod_copy_async(&src.bias_sigma, &mut self.bias_sigma, "bias_sigma", &self.stream)?;
Self::dtod_copy_async(&src.weight_epsilon, &mut self.weight_epsilon, "weight_epsilon", &self.stream)?;
Self::dtod_copy_async(&src.bias_epsilon, &mut self.bias_epsilon, "bias_epsilon", &self.stream)?;
// Task 2.5 Bug #4: propagate scheduler sigma scale so a target-net
// sync doesn't reset the effective σ reporting to 1.0.
self.current_sigma_scale = src.current_sigma_scale;
Ok(())
}