Merge pull request #56 from arvigj/stochastic_gd

Add stochastic gradient descent
polyfem · Dec 8, 2023 · e9bb718 · e9bb718
2 parents 2c79268 + 1d16421
commit e9bb718
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 3 deletions.
diff --git a/non-linear-solver-spec.json b/non-linear-solver-spec.json
@@ -15,6 +15,7 @@
             "LBFGS",
             "LBFGSB",
             "Newton",
+            "StochasticGradientDescent",
             "box_constraints",
             "advanced"
         ],
@@ -28,6 +29,7 @@
             "Newton",
             "DenseNewton",
             "GradientDescent",
+            "StochasticGradientDescent",
             "L-BFGS",
             "BFGS",
             "L-BFGS-B",
@@ -164,6 +166,21 @@
         "type": "bool",
         "doc": "Use PSD as fallback using second order solvers (i.e., Newton's method)."
     },
+    {
+        "pointer": "/StochasticGradientDescent",
+        "default": null,
+        "type": "object",
+        "optional": [
+            "erase_component_probability"
+        ],
+        "doc": "Options for Stochastic Gradient Descent."
+    },
+    {
+        "pointer": "/StochasticGradientDescent/erase_component_probability",
+        "default": 0.3,
+        "type": "float",
+        "doc": "Probability of erasing a component on the gradient for StochasticGradientDescent."
+    },
     {
         "pointer": "/line_search",
         "default": null,

diff --git a/src/polysolve/nonlinear/Solver.cpp b/src/polysolve/nonlinear/Solver.cpp
@@ -85,6 +85,11 @@ namespace polysolve::nonlinear
             solver->add_strategy(std::make_unique<LBFGS>(
                 solver_params, characteristic_length, logger));
         }
+        else if (solver_name == "StochasticGradientDescent" || solver_name == "stochastic_gradient_descent")
+        {
+            solver->add_strategy(std::make_unique<GradientDescent>(
+                solver_params, true, characteristic_length, logger));
+        }
         else if (solver_name == "GradientDescent" || solver_name == "gradient_descent")
         {
             // grad descent always there
@@ -93,7 +98,7 @@ namespace polysolve::nonlinear
             throw std::runtime_error("Unrecognized solver type: " + solver_name);
 
         solver->add_strategy(std::make_unique<GradientDescent>(
-            solver_params, characteristic_length, logger));
+            solver_params, false, characteristic_length, logger));
 
         solver->set_strategies_iterations(solver_params);
         return solver;
@@ -105,6 +110,7 @@ namespace polysolve::nonlinear
                 "DenseNewton",
                 "Newton",
                 "GradientDescent",
+                "StochasticGradientDescent",
                 "L-BFGS"};
     }
 

diff --git a/src/polysolve/nonlinear/descent_strategies/GradientDescent.cpp b/src/polysolve/nonlinear/descent_strategies/GradientDescent.cpp
@@ -4,10 +4,13 @@ namespace polysolve::nonlinear
 {
 
     GradientDescent::GradientDescent(const json &solver_params_,
+                                     const bool is_stochastic,
                                      const double characteristic_length,
                                      spdlog::logger &logger)
-        : Superclass(solver_params_, characteristic_length, logger)
+        : Superclass(solver_params_, characteristic_length, logger), is_stochastic_(is_stochastic)
     {
+        if (is_stochastic_)
+            erase_component_probability_ = solver_params_["StochasticGradientDescent"]["erase_component_probability"];
     }
 
     bool GradientDescent::compute_update_direction(
@@ -18,6 +21,13 @@ namespace polysolve::nonlinear
     {
         direction = -grad;
 
+        if (is_stochastic_)
+        {
+            Eigen::VectorXd mask = (Eigen::VectorXd::Random(direction.size()).array() + 1.) / 2.;
+            for (int i = 0; i < direction.size(); ++i)
+                direction(i) *= (mask(i) < erase_component_probability_) ? 0. : 1.;
+        }
+
         return true;
     }
 

diff --git a/src/polysolve/nonlinear/descent_strategies/GradientDescent.hpp b/src/polysolve/nonlinear/descent_strategies/GradientDescent.hpp
@@ -13,15 +13,20 @@ namespace polysolve::nonlinear
         using Superclass = DescentStrategy;
 
         GradientDescent(const json &solver_params_,
+                        const bool is_stochastic,
                         const double characteristic_length,
                         spdlog::logger &logger);
 
-        std::string name() const override { return "GradientDescent"; }
+        std::string name() const override { return is_stochastic_ ? "StochasticGradientDescent" : "GradientDescent"; }
 
         bool compute_update_direction(
             Problem &objFunc,
             const TVector &x,
             const TVector &grad,
             TVector &direction) override;
+
+    private:
+        bool is_stochastic_ = false;
+        double erase_component_probability_ = 0;
     };
 } // namespace polysolve::nonlinear