From 0bd5c68f56e397e5ea0e9c0846743881fbc44943 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 20 Apr 2026 21:25:55 +0200 Subject: [PATCH] feat(F3): spectral_gap via rolling Gram + power iteration (sigma_1/sigma_2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace single-sample max/min proxy with mathematically proper sigma_1/sigma_2 ratio. Maintains a host-side VecDeque of up to 64 Q-value samples; once ≥ 8 samples are available computes the n_cols×n_cols Gram matrix X^T X, then extracts the two largest singular values via power iteration + rank-one deflation. Falls back to the coarse max/min ratio until the buffer fills. compute_q_spectral_gap changed to &mut self; callers updated accordingly. Co-Authored-By: Claude Sonnet 4.6 --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 139 ++++++++++++++++-- crates/ml/src/trainers/dqn/fused_training.rs | 4 +- .../src/trainers/dqn/trainer/training_loop.rs | 4 +- 3 files changed, 128 insertions(+), 19 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 72e83e9c0..c813cc529 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1685,6 +1685,11 @@ pub struct GpuDqnTrainer { pub(crate) q_snapshots: crate::cuda_pipeline::q_snapshot::SnapshotRing, /// D1/N1: Whether distillation gradient was applied in the most recent epoch boundary. pub(crate) last_distill_active: bool, + + /// F3: Rolling buffer of the last up-to-64 Q-value samples (each `total_actions` floats) + /// read from q_readback_pinned[7..]. Used by `compute_q_spectral_gap` to estimate + /// sigma_1/sigma_2 via Gram power iteration. Capped at 64 samples × 12 floats = 3 KB. + pub(crate) q_sample_history: std::collections::VecDeque>, } impl GpuDqnTrainer { @@ -7482,6 +7487,7 @@ impl GpuDqnTrainer { total_params, ), last_distill_active: false, + q_sample_history: std::collections::VecDeque::with_capacity(64), }) } @@ -7872,37 +7878,84 @@ impl GpuDqnTrainer { Ok(out) } - /// Coarse Q-collapse detector derived from the single-sample Q readback at - /// `q_readback_pinned[7..7+total_actions]`. + /// F3: Spectral gap — sigma_1 / sigma_2 of the rolling Q-sample matrix. + /// High ratio = rank-1 collapse (one direction dominates); ~1.0 = balanced. /// - /// Returns a LARGE value (~100) when all Q values are (nearly) equal — i.e. - /// rank-1 / Q-collapse — and a value close to 1.0 when the Q spread is wide. - /// Intentionally coarse: a full SVD across a batch of Q samples would require - /// a separate larger buffer and is out of scope for the host-side health path. + /// Maintains a host-side rolling buffer of the last 64 Q-value samples + /// (each `total_actions` floats from `q_readback_pinned[7..]`). Once ≥ 8 + /// samples are available, computes sigma_1/sigma_2 via power iteration on + /// the n_cols × n_cols Gram matrix X^T X (matrix is tiny — instant on host). + /// Falls back to single-sample max/min ratio when the buffer is too small. /// /// The downstream `NormalizedComponents::from_raw` thresholds are calibrated /// for this behaviour (`spectral_gap_norm` drops toward 0 as this value rises). - pub fn compute_q_spectral_gap(&self) -> f32 { + pub fn compute_q_spectral_gap(&mut self) -> f32 { if self.q_readback_pinned.is_null() { return 1.0; } - // Layout: q_readback_pinned[0..7] = stats, [7..7+total_actions] = per-action Q values. - let total_actions = self.total_actions() as usize; - if total_actions == 0 { return 1.0; } + + // Pull the latest sample (total_actions floats) from the readback head. + let n_cols = self.total_actions(); + if n_cols == 0 { return 1.0; } + let mut sample = Vec::with_capacity(n_cols); unsafe { let base = self.q_readback_pinned.add(7); + for i in 0..n_cols { + let v = *base.add(i); + if !v.is_finite() { return 1.0; } + sample.push(v); + } + } + + // Push into rolling history (capacity 64). + if self.q_sample_history.len() >= 64 { self.q_sample_history.pop_front(); } + self.q_sample_history.push_back(sample); + + if self.q_sample_history.len() < 8 { + // Not enough samples yet — fall back to coarse max/min ratio of the + // latest sample (pre-F3 behaviour). + let latest = self.q_sample_history.back().unwrap(); let mut max = f32::NEG_INFINITY; let mut min = f32::INFINITY; - for i in 0..total_actions { - let v = *base.add(i); + for &v in latest { if v.is_finite() { if v > max { max = v; } if v < min { min = v; } } } let range = (max - min).abs(); - // Range near zero means all actions have equal Q → rank-1-like collapse; - // return a large value to signal pathology via spectral_gap_norm ≈ 0. - if range < 1e-6 { 100.0 } else { 1.0 + range.recip() } + return if range < 1e-6 { 100.0 } else { 1.0 + range.recip() }; } + + // Build centered X matrix (B × n_cols), compute column means first. + let b = self.q_sample_history.len(); + let mut means = vec![0.0f32; n_cols]; + for s in &self.q_sample_history { + for j in 0..n_cols { means[j] += s[j]; } + } + for m in &mut means { *m /= b as f32; } + + // Gram matrix G = X^T X, an n_cols × n_cols symmetric PSD matrix. + let mut gram = vec![0.0f32; n_cols * n_cols]; + for s in &self.q_sample_history { + for i in 0..n_cols { + let ci = s[i] - means[i]; + for j in 0..n_cols { + let cj = s[j] - means[j]; + gram[i * n_cols + j] += ci * cj; + } + } + } + + // Power iteration for largest eigenvalue (lambda_1 = sigma_1^2). + let lambda_1 = power_iteration_largest(&gram, n_cols, 40); + // Deflate and compute second-largest. + let deflated = deflate_rank_one(&gram, n_cols, lambda_1); + let lambda_2 = power_iteration_largest(&deflated, n_cols, 40); + + if lambda_2 < 1e-12 { return 100.0; } + let sigma1 = lambda_1.max(0.0).sqrt(); + let sigma2 = lambda_2.max(0.0).sqrt(); + if sigma2 < 1e-6 { return 100.0; } + sigma1 / sigma2 } /// Read eval_q_std_ema. @@ -12570,6 +12623,62 @@ impl GpuDqnTrainer { } } +// ── F3: Spectral-gap power-iteration helpers ───────────────────────────────── + +/// F3: Largest eigenvalue of a symmetric matrix via power iteration. +/// `mat` is stored row-major, size `n × n`. Runs `max_iters` iterations. +/// Numerically stable: renormalizes each step. +fn power_iteration_largest(mat: &[f32], n: usize, max_iters: usize) -> f32 { + if n == 0 { return 0.0; } + let mut v = vec![1.0f32 / (n as f32).sqrt(); n]; + let mut lambda = 0.0f32; + for _ in 0..max_iters { + let mut av = vec![0.0f32; n]; + for i in 0..n { + for j in 0..n { av[i] += mat[i * n + j] * v[j]; } + } + let norm = av.iter().map(|x| x * x).sum::().sqrt(); + if norm < 1e-20 { return 0.0; } + let inv = 1.0 / norm; + for i in 0..n { v[i] = av[i] * inv; } + // Rayleigh quotient + let mut new_lambda = 0.0f32; + let mut mv = vec![0.0f32; n]; + for i in 0..n { + for j in 0..n { mv[i] += mat[i * n + j] * v[j]; } + } + for i in 0..n { new_lambda += v[i] * mv[i]; } + if (new_lambda - lambda).abs() < 1e-9 { return new_lambda; } + lambda = new_lambda; + } + lambda +} + +/// F3: Deflate a symmetric matrix by removing its dominant eigenpair. +/// Given largest eigenvalue `lambda_1` (approximately), returns mat - lambda_1 * v_1 v_1^T +/// where v_1 is extracted via one more power iteration. +fn deflate_rank_one(mat: &[f32], n: usize, lambda_1: f32) -> Vec { + // Recover v_1 by one more round of power iteration (cheap, n=12). + let mut v = vec![1.0f32 / (n as f32).sqrt(); n]; + for _ in 0..40 { + let mut av = vec![0.0f32; n]; + for i in 0..n { + for j in 0..n { av[i] += mat[i * n + j] * v[j]; } + } + let norm = av.iter().map(|x| x * x).sum::().sqrt(); + if norm < 1e-20 { break; } + let inv = 1.0 / norm; + for i in 0..n { v[i] = av[i] * inv; } + } + let mut out = mat.to_vec(); + for i in 0..n { + for j in 0..n { + out[i * n + j] -= lambda_1 * v[i] * v[j]; + } + } + out +} + // ── Compilation ───────────────────────────────────────────────────────────── /// Load the utility kernels from precompiled cubin (ZERO runtime nvcc). diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index e62c95355..acee01032 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -2249,8 +2249,8 @@ impl FusedTrainingCtx { /// B4/G5: Last ensemble gradient budget (constant 0.05). pub(crate) fn last_ens_budget_eff(&self) -> f32 { self.trainer.last_ens_budget_eff } - /// Compute spectral gap on the last Q readback slice (coarse max/min ratio proxy). - pub(crate) fn compute_q_spectral_gap(&self) -> f32 { + /// F3: Spectral gap — sigma_1 / sigma_2 via rolling Gram + power iteration. + pub(crate) fn compute_q_spectral_gap(&mut self) -> f32 { self.trainer.compute_q_spectral_gap() } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 671f60000..0f91a5427 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1815,8 +1815,8 @@ impl DQNTrainer { .map(|f| f.read_atom_utilization()) .unwrap_or(0.0); - // Spectral gap proxy from q_readback_pinned[7..7+total_actions]. - let spectral_gap = self.fused_ctx.as_ref() + // F3: Spectral gap — sigma_1/sigma_2 via rolling Gram + power iteration. + let spectral_gap = self.fused_ctx.as_mut() .map(|f| f.compute_q_spectral_gap()) .unwrap_or(1.0);