From 3af3eeba59cfcbc16acbab215addd2bbd3cf2528 Mon Sep 17 00:00:00 2001 From: sharmis-tp Date: Tue, 10 Nov 2026 02:06:20 +0000 Subject: [PATCH 1/2] fix centered RMSprop: clamp variance, move eps outside sqrt --- src/optim/rmsprop.rs | 49 +++++++++++++++++++++++++++++++------------- 1 file changed, 35 insertions(+), 14 deletions(-) diff --git a/src/optim/rmsprop.rs b/src/optim/rmsprop.rs index 60af625..b7ec39f 100644 --- a/src/optim/rmsprop.rs +++ b/src/optim/rmsprop.rs @@ -84,12 +84,11 @@ impl Optimizer for RMSprop { let square = blend(self.square_avg[index].as_ref(), &grad.square(), self.alpha); - // The variance can dip below zero in floating point, so eps goes inside the root. let denominator = if self.centered { let mean = blend(self.grad_avg[index].as_ref(), &grad, self.alpha); let variance = square.sub(&mean.square()); self.grad_avg[index] = Some(mean); - variance.add_scalar(self.eps).sqrt() + variance.clamp(0.0, f32::INFINITY).sqrt().add_scalar(self.eps) } else { square.sqrt().add_scalar(self.eps) }; @@ -160,31 +159,53 @@ mod tests { ); } - /// With a constant gradient the variance estimate shrinks, so late steps - /// dwarf early ones. The plain form settles at |g| instead. + /// With a constant gradient, centered RMSprop drives variance to zero, + /// so the denominator converges to eps. Step size should settle to a + /// steady value -- not keep growing without bound. #[test] - fn centered_variant_accelerates_on_constant_gradients() { + fn centered_variant_stabilises_on_constant_gradients() { let param = Param::new(Tensor::from_vec(vec![0.0], &[1])); let mut opt = RMSprop::new(vec![param.clone()], 0.01) .alpha(0.9) .centered(); let mut previous = 0.0f32; - let mut first_step_size = 0.0f32; - let mut step_size = 0.0f32; - for iteration in 0..100 { + let mut step_sizes = Vec::new(); + for _ in 0..200 { param.set_grad(Tensor::from_vec(vec![1.0], &[1])); opt.step(); let value = param.value().to_vec()[0]; - step_size = (previous - value).abs(); - if iteration == 0 { - first_step_size = step_size; - } + step_sizes.push((previous - value).abs()); previous = value; } + // The last few steps should have converged: within 1% of each other. + let last = step_sizes[step_sizes.len() - 1]; + let prev = step_sizes[step_sizes.len() - 2]; + assert!( + (last - prev).abs() / last < 0.01, + "step size should stabilise, but last two were {prev} and {last}" + ); + // And none should be NaN or infinite. assert!( - step_size > 5.0 * first_step_size, - "late steps ({step_size}) should dwarf the first ({first_step_size})" + step_sizes.iter().all(|s| s.is_finite()), + "step sizes must stay finite" ); } + + /// Large gradients (~1e4) trigger cancellation error in E[g^2] - E[g]^2. + /// The variance can go negative; verify we clamp it and never produce NaN. + #[test] + fn centered_large_gradients_no_nan() { + let param = Param::new(Tensor::from_vec(vec![0.0], &[1])); + let mut opt = RMSprop::new(vec![param.clone()], 1e-3) + .alpha(0.99) + .centered(); + + for _ in 0..100 { + param.set_grad(Tensor::from_vec(vec![1e4], &[1])); + opt.step(); + let value = param.value().to_vec()[0]; + assert!(value.is_finite(), "parameter became NaN or Inf"); + } + } } From 32fab00ded9cfea892f0cf482d5130155d440099 Mon Sep 17 00:00:00 2001 From: Canberk Pitirli Date: Thu, 12 Nov 2026 13:11:19 +0300 Subject: [PATCH 2/2] keep eps inside the root and bound the step in the tests --- src/optim/rmsprop.rs | 57 ++++++++++++++++++++++---------------------- 1 file changed, 28 insertions(+), 29 deletions(-) diff --git a/src/optim/rmsprop.rs b/src/optim/rmsprop.rs index b7ec39f..a05361d 100644 --- a/src/optim/rmsprop.rs +++ b/src/optim/rmsprop.rs @@ -88,7 +88,10 @@ impl Optimizer for RMSprop { let mean = blend(self.grad_avg[index].as_ref(), &grad, self.alpha); let variance = square.sub(&mean.square()); self.grad_avg[index] = Some(mean); - variance.clamp(0.0, f32::INFINITY).sqrt().add_scalar(self.eps) + // E[g²] - E[g]² cancels for large gradients and can come out + // negative (#12), and eps stays inside the root so the step + // never exceeds lr·|g|/√eps when the variance estimate vanishes. + variance.relu().add_scalar(self.eps).sqrt() } else { square.sqrt().add_scalar(self.eps) }; @@ -159,53 +162,49 @@ mod tests { ); } - /// With a constant gradient, centered RMSprop drives variance to zero, - /// so the denominator converges to eps. Step size should settle to a - /// steady value -- not keep growing without bound. + /// A constant gradient drives the variance estimate to exactly zero, so + /// the denominator bottoms out at √eps and the step at lr·|g|/√eps. #[test] - fn centered_variant_stabilises_on_constant_gradients() { + fn centered_variant_caps_the_step_on_constant_gradients() { + let lr = 0.01; let param = Param::new(Tensor::from_vec(vec![0.0], &[1])); - let mut opt = RMSprop::new(vec![param.clone()], 0.01) - .alpha(0.9) - .centered(); + let mut opt = RMSprop::new(vec![param.clone()], lr).alpha(0.9).centered(); + let cap = lr / 1e-8f32.sqrt(); let mut previous = 0.0f32; - let mut step_sizes = Vec::new(); for _ in 0..200 { param.set_grad(Tensor::from_vec(vec![1.0], &[1])); opt.step(); let value = param.value().to_vec()[0]; - step_sizes.push((previous - value).abs()); + let step = (previous - value).abs(); + assert!( + step.is_finite() && step <= cap * 1.01, + "step {step} past {cap}" + ); previous = value; } - // The last few steps should have converged: within 1% of each other. - let last = step_sizes[step_sizes.len() - 1]; - let prev = step_sizes[step_sizes.len() - 2]; - assert!( - (last - prev).abs() / last < 0.01, - "step size should stabilise, but last two were {prev} and {last}" - ); - // And none should be NaN or infinite. - assert!( - step_sizes.iter().all(|s| s.is_finite()), - "step sizes must stay finite" - ); } - /// Large gradients (~1e4) trigger cancellation error in E[g^2] - E[g]^2. - /// The variance can go negative; verify we clamp it and never produce NaN. + /// Gradients around 1e4 make E[g²] - E[g]² cancel to a negative value; + /// the clamp keeps the denominator real and the step bounded. #[test] - fn centered_large_gradients_no_nan() { + fn centered_large_gradients_stay_finite_and_bounded() { + let lr = 1e-3; let param = Param::new(Tensor::from_vec(vec![0.0], &[1])); - let mut opt = RMSprop::new(vec![param.clone()], 1e-3) - .alpha(0.99) - .centered(); + let mut opt = RMSprop::new(vec![param.clone()], lr).alpha(0.99).centered(); + let cap = lr * 1e4 / 1e-8f32.sqrt(); + let mut previous = 0.0f32; for _ in 0..100 { param.set_grad(Tensor::from_vec(vec![1e4], &[1])); opt.step(); let value = param.value().to_vec()[0]; - assert!(value.is_finite(), "parameter became NaN or Inf"); + let step = (previous - value).abs(); + assert!( + step.is_finite() && step <= cap * 1.01, + "step {step} past {cap}" + ); + previous = value; } } }