Fix centered RMSprop: clamp the variance, keep eps inside the root

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})"
-        );
     }
 }