Normalize step size to the batch size used (#431)

* Normalize step size to the batch size used.

* Update history.

* Don't use if constexpr.

* Fix all step sizes in tests to account for new behavior.

* Retune DemonSGD tests to reduce failure probability.

* Fix merge artifact.

* Fix MOEAD test to be more robust.

Author: rcurtin

HISTORY.md

@@ -1,5 +1,10 @@
 ### ensmallen ?.??.?: "???"
 ###### ????-??-??
+ * SGD-like optimizers now all divide the step size by the batch size so that
+   step sizes don't need to be tuned in addition to batch sizes.  If you require
+   behavior from ensmallen 2, define the `ENS_OLD_SEPARABLE_STEP_BEHAVIOR` macro
+   before including `ensmallen.hpp`
+   ([#431](https://github.com/mlpack/ensmallen/pull/431)).
 
 ### ensmallen 2.22.2: "E-Bike Excitement"
 ###### 2025-04-30

include/ensmallen_bits/cmaes/active_cmaes_impl.hpp

@@ -110,7 +110,7 @@ typename MatType::elem_type ActiveCMAES<SelectionPolicyType,
   BaseMatType sigma(2, 1); // sigma is vector-shaped.
   if (stepSize == 0)
     sigma(0) = transformationPolicy.InitialStepSize();
-  else 
+  else
     sigma(0) = ElemType(stepSize);
 
   const ElemType cs = 4 / ElemType(iterate.n_elem + 4);

include/ensmallen_bits/eve/eve_impl.hpp

@@ -149,8 +149,14 @@ Eve::Optimize(SeparableFunctionType& function,
 
     lastObjective = objective;
 
+    // TODO: remove in ensmallen 4.0.0.
+    #if defined(ENS_OLD_SEPARABLE_STEP_BEHAVIOR)
     iterate -= ElemType(stepSize) / dt * (m / biasCorrection1) /
         (sqrt(v / biasCorrection2) + ElemType(epsilon));
+    #else
+    iterate -= (ElemType(stepSize) / (dt * effectiveBatchSize)) *
+        (m / biasCorrection1) / (sqrt(v / biasCorrection2) + ElemType(epsilon));
+    #endif
 
     terminate |= Callback::StepTaken(*this, f, iterate, callbacks...);
 

include/ensmallen_bits/lookahead/lookahead_impl.hpp

@@ -66,6 +66,24 @@ inline Lookahead<BaseOptimizerType, DecayPolicyType>::~Lookahead()
   instDecayPolicy.Clean();
 }
 
+template<typename BaseOptimizerType>
+size_t GetBatchSize(
+    const BaseOptimizerType& baseOptimizer,
+    const typename std::enable_if_t<traits::HasBatchSizeSignature<
+        BaseOptimizerType>::value>* = 0)
+{
+  return baseOptimizer.BatchSize();
+}
+
+template<typename BaseOptimizerType>
+size_t GetBatchSize(
+    const BaseOptimizerType& baseOptimizer,
+    const typename std::enable_if_t<!traits::HasBatchSizeSignature<
+        BaseOptimizerType>::value>* = 0)
+{
+  return 1;
+}
+
 //! Optimize the function (minimize).
 template<typename BaseOptimizerType, typename DecayPolicyType>
 template<typename SeparableFunctionType,
@@ -189,11 +207,9 @@ Lookahead<BaseOptimizerType, DecayPolicyType>::Optimize(
     // Find the number of functions to use.
     const size_t numFunctions = f.NumFunctions();
 
-    size_t batchSize = 1;
     // Check if the optimizer implements the BatchSize() method and use the
     // parameter for the objective calculation.
-    if (traits::HasBatchSizeSignature<BaseOptimizerType>::value)
-      batchSize = baseOptimizer.BatchSize();
+    size_t batchSize = GetBatchSize(baseOptimizer);
 
     overallObjective = 0;
     for (size_t i = 0; i < numFunctions; i += batchSize)

include/ensmallen_bits/parallel_sgd/parallel_sgd_impl.hpp

@@ -132,7 +132,7 @@ typename MatType::elem_type>::type ParallelSGD<DecayPolicyType>::Optimize(
       return overallObjective;
     }
 
-    // Get the stepsize for this iteration
+    // Get the stepsize for this iteration.
     double stepSize = decayPolicy.StepSize(i);
 
     // Shuffle for uniform sampling of functions by each thread.
@@ -180,6 +180,8 @@ typename MatType::elem_type>::type ParallelSGD<DecayPolicyType>::Optimize(
 
             // Call out to utility function to use the right type of OpenMP
             // lock.
+            // TODO: if batch size support > 1 is added, `stepSize` will need to
+            // be updated here.
             UpdateLocation(iterate, row, i, ElemType(stepSize) * value);
           }
         }
@@ -198,4 +200,4 @@ typename MatType::elem_type>::type ParallelSGD<DecayPolicyType>::Optimize(
 
 } // namespace ens
 
-#endif
+#endif

include/ensmallen_bits/sgd/sgd_impl.hpp

@@ -150,8 +150,14 @@ SGD<UpdatePolicyType, DecayPolicyType>::Optimize(
         gradient, callbacks...);
 
     // Use the update policy to take a step.
+    // TODO: remove old behavior in ensmallen 4.0.0.
+    #if defined(ENS_OLD_SEPARABLE_STEP_BEHAVIOR)
     instUpdatePolicy.As<InstUpdatePolicyType>().Update(iterate, stepSize,
         gradient);
+    #else
+    instUpdatePolicy.As<InstUpdatePolicyType>().Update(iterate,
+        (stepSize / effectiveBatchSize), gradient);
+    #endif
 
     terminate |= Callback::StepTaken(*this, f, iterate, callbacks...);
 

tests/ada_belief_test.cpp

@@ -22,7 +22,7 @@ using namespace ens::test;
 TEMPLATE_TEST_CASE("AdaBelief_LogisticRegressionFunction", "[AdaBelief]",
     ENS_ALL_TEST_TYPES)
 {
-  AdaBelief adaBelief;
+  AdaBelief adaBelief(0.032);
   LogisticRegressionFunctionTest<TestType, arma::Row<size_t>>(adaBelief);
 }
 
@@ -31,7 +31,7 @@ TEMPLATE_TEST_CASE("AdaBelief_LogisticRegressionFunction", "[AdaBelief]",
 TEMPLATE_TEST_CASE("AdaBelief_LogisticRegressionFunction", "[AdaBelief]",
     coot::mat, coot::fmat)
 {
-  AdaBelief adaBelief;
+  AdaBelief adaBelief(0.032);
   LogisticRegressionFunctionTest<TestType, coot::Row<size_t>>(adaBelief);
 }
 

tests/ada_bound_test.cpp

@@ -22,7 +22,7 @@ using namespace ens::test;
 TEMPLATE_TEST_CASE("AdaBound_SphereFunction", "[AdaBound]", ENS_ALL_TEST_TYPES,
     ENS_SPARSE_TEST_TYPES)
 {
-  AdaBound optimizer(0.005, 2, 0.1, 1e-3, 0.9, 0.999, 1e-8, 500000,
+  AdaBound optimizer(0.002, 2, 0.1, 1e-3, 0.9, 0.999, 1e-8, 500000,
       1e-3, false);
   FunctionTest<SphereFunction, TestType>(
       optimizer,
@@ -33,7 +33,7 @@ TEMPLATE_TEST_CASE("AdaBound_SphereFunction", "[AdaBound]", ENS_ALL_TEST_TYPES,
 TEMPLATE_TEST_CASE("AMSBound_SphereFunction", "[AdaBound]", ENS_ALL_TEST_TYPES,
     ENS_SPARSE_TEST_TYPES)
 {
-  AMSBound optimizer(0.005, 2, 0.1, 1e-3, 0.9, 0.999, 1e-8, 500000,
+  AMSBound optimizer(0.002, 2, 0.1, 1e-3, 0.9, 0.999, 1e-8, 500000,
       1e-3, false);
   FunctionTest<SphereFunction, TestType>(
       optimizer,
@@ -47,7 +47,7 @@ TEMPLATE_TEST_CASE("AdaBound_SphereFunctionSpMatDenseGradient", "[AdaBound]",
   typedef typename TestType::elem_type ElemType;
 
   SphereFunction f(2);
-  AdaBound optimizer(0.001, 2, 0.1, 1e-3, 0.9, 0.999, 1e-8, 500000,
+  AdaBound optimizer(0.002, 2, 0.1, 1e-3, 0.9, 0.999, 1e-8, 500000,
       1e-3, false);
 
   TestType coordinates = arma::conv_to<TestType>::from(
@@ -65,7 +65,7 @@ TEMPLATE_TEST_CASE("AMSBound_SphereFunctionSpMatDenseGradient", "[AdaBound]",
   typedef typename TestType::elem_type ElemType;
 
   SphereFunction f(2);
-  AMSBound optimizer(0.001, 2, 0.1, 1e-3, 0.9, 0.999, 1e-8, 500000,
+  AMSBound optimizer(0.002, 2, 0.1, 1e-3, 0.9, 0.999, 1e-8, 500000,
       1e-3, false);
 
   arma::sp_mat coordinates = f.GetInitialPoint<TestType>();
@@ -81,15 +81,15 @@ TEMPLATE_TEST_CASE("AMSBound_SphereFunctionSpMatDenseGradient", "[AdaBound]",
 TEMPLATE_TEST_CASE("AdaBound_SphereFunction", "[AdaBound]",
     coot::mat, coot::fmat)
 {
-  AdaBound optimizer(0.001, 2, 0.1, 1e-3, 0.9, 0.999, 1e-8, 500000,
+  AdaBound optimizer(0.002, 2, 0.1, 1e-3, 0.9, 0.999, 1e-8, 500000,
       1e-3, false);
   FunctionTest<SphereFunction, TestType>(optimizer, 0.5, 0.1);
 }
 
 TEMPLATE_TEST_CASE("AMSBoundSphereFunctionTest", "[AdaBound]",
     coot::mat, coot::fmat)
 {
-  AMSBound optimizer(0.001, 2, 0.1, 1e-3, 0.9, 0.999, 1e-8, 500000,
+  AMSBound optimizer(0.002, 2, 0.1, 1e-3, 0.9, 0.999, 1e-8, 500000,
       1e-3, false);
   FunctionTest<SphereFunction, TestType>(optimizer, 0.5, 0.1);
 }

tests/ada_delta_test.cpp

@@ -29,7 +29,7 @@ TEMPLATE_TEST_CASE("AdaDelta_LogisticRegressionFunction", "[AdaDelta]",
 
   // Use a large epsilon if we are using FP16, to avoid underflow in the first
   // iterations.
-  AdaDelta adaDelta(1.0, 32, 0.95, sizeof(ElemType) == 2 ? 1e-4 : 1e-6);
+  AdaDelta adaDelta(32.0, 32, 0.95, sizeof(ElemType) == 2 ? 1e-4 : 1e-6);
   LogisticRegressionFunctionTest<TestType, arma::Row<size_t>>(adaDelta);
 }
 
@@ -38,7 +38,7 @@ TEMPLATE_TEST_CASE("AdaDelta_LogisticRegressionFunction", "[AdaDelta]",
 TEMPLATE_TEST_CASE("AdaDelta_LogisticRegressionFunction", "[AdaDelta]",
     coot::mat, coot::fmat)
 {
-  AdaDelta adaDelta;
+  AdaDelta adaDelta(32.0);
   LogisticRegressionFunctionTest<TestType, coot::Row<size_t>>(
       adaDelta, 0.003, 0.006, 1);
 }

tests/ada_grad_test.cpp

@@ -23,7 +23,7 @@ using namespace ens::test;
 TEMPLATE_TEST_CASE("AdaGrad_LogisticRegressionFunction", "[AdaGrad]",
     ENS_ALL_TEST_TYPES)
 {
-  AdaGrad adagrad(0.99, 32, 10 * Tolerances<TestType>::Obj, 500000,
+  AdaGrad adagrad(31.5, 32, 10 * Tolerances<TestType>::Obj, 500000,
       Tolerances<TestType>::Obj, true);
   LogisticRegressionFunctionTest<TestType, arma::Row<size_t>>(adagrad);
 }
@@ -33,7 +33,7 @@ TEMPLATE_TEST_CASE("AdaGrad_LogisticRegressionFunction", "[AdaGrad]",
 TEMPLATE_TEST_CASE("AdaGrad_LogisticRegressionFunction", "[AdaGrad]",
     coot::mat, coot::fmat)
 {
-  AdaGrad adagrad(0.99, 32, 1e-8, 500000, 1e-9, true);
+  AdaGrad adagrad(31.5, 32, 1e-8, 500000, 1e-9, true);
   LogisticRegressionFunctionTest<TestType, coot::Row<size_t>>(adagrad);
 }