diff --git a/src/optim/rmsprop.rs b/src/optim/rmsprop.rs index 60af625..a05361d 100644 --- a/src/optim/rmsprop.rs +++ b/src/optim/rmsprop.rs @@ -84,12 +84,14 @@ 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() + // 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) }; @@ -160,31 +162,49 @@ mod tests { ); } - /// With a constant gradient the variance estimate shrinks, so late steps - /// dwarf early ones. The plain form settles at |g| instead. + /// 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_accelerates_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 first_step_size = 0.0f32; - let mut step_size = 0.0f32; - for iteration in 0..100 { + 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; - } + let step = (previous - value).abs(); + assert!( + step.is_finite() && step <= cap * 1.01, + "step {step} past {cap}" + ); + previous = value; + } + } + + /// 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_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()], 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]; + let step = (previous - value).abs(); + assert!( + step.is_finite() && step <= cap * 1.01, + "step {step} past {cap}" + ); previous = value; } - assert!( - step_size > 5.0 * first_step_size, - "late steps ({step_size}) should dwarf the first ({first_step_size})" - ); } }