Optimizer: cap the first L-BFGS/BFGS step so that large initial gradients do not fail the first line search

src/stan/optimization/bfgs.hpp

@@ -48,6 +48,10 @@ class LSOptions {
   Scalar minAlpha{1e-12};
   Scalar maxLSIts{20};
   Scalar maxLSRestarts{10};
+  // largest change of any parameter in the first trial step of the first
+  // iteration and after a Hessian reset (the step is alpha0 times the
+  // gradient, or shorter)
+  Scalar maxStep0{std::numeric_limits<Scalar>::infinity()};
 };
 template <typename FunctorType, typename QNUpdateType, typename Scalar = double,
           int DimAtCompile = Eigen::Dynamic>
@@ -62,6 +66,7 @@ class BFGSMinimizer {
   Scalar _fk, _fk_1, _alphak_1;
   Scalar _alpha, _alpha0;
   size_t _itNum;
+  bool _capped{false};
   std::string _note;
   QNUpdateType _qn;
 
@@ -159,6 +164,7 @@ class BFGSMinimizer {
     _pk = -_gk;
 
     _itNum = 0;
+    _capped = false;
     _note = "";
   }
 
@@ -170,9 +176,14 @@ class BFGSMinimizer {
 
     _itNum++;
 
-    if (_itNum == 1) {
+    // If maxStep0 cut the first step short, start once more from -g: the
+    // pair (s, y) of that step measures only the parameters with the largest
+    // gradient, and an approximation built from it makes the next steps in
+    // all other parameters tiny as well. Restarting after every capped step
+    // instead turns runs with large gradients into steepest descent.
+    if (_itNum == 1 || _capped) {
       resetB = 1;
-      _note = "";
+      _note = _capped ? "capped first step, Hessian reset" : "";
     } else {
       resetB = 0;
       _note = "";
@@ -185,23 +196,30 @@ class BFGSMinimizer {
       }
 
       // Get an initial guess for the step size (alpha)
-      if (_itNum > 1 && resetB != 2) {
+      if (!resetB) {
         // use cubic interpolation based on the previous step
         _alpha0 = _alpha = std::min(
             1.0, 1.01
                      * CubicInterp(_gk_1.dot(_pk_1), _alphak_1, _fk - _fk_1,
                                    _gk.dot(_pk_1), _ls_opts.minAlpha, 1.0));
       } else {
-        // On the first step (or, after a reset) use the default step size
-        _alpha0 = _alpha = _ls_opts.alpha0;
+        // On the first step (or, after a reset) use the default step size,
+        // but move no parameter by more than maxStep0 along the gradient
+        _alpha0 = _alpha = std::min(
+            _ls_opts.alpha0,
+            _ls_opts.maxStep0 / _pk.template lpNorm<Eigen::Infinity>());
       }
 
       // Perform the line search.  If successful, the results are in the
-      // variables: _xk_1, _fk_1 and _gk_1.
-      retCode
-          = WolfeLineSearch(_func, _alpha, _xk_1, _fk_1, _gk_1, _pk, _xk, _fk,
-                            _gk, _ls_opts.c1, _ls_opts.c2, _ls_opts.minAlpha,
-                            _ls_opts.maxLSIts, _ls_opts.maxLSRestarts);
+      // variables: _xk_1, _fk_1 and _gk_1. minAlpha is the smallest step the
+      // line search may return. A first trial step that maxStep0 shortened
+      // gets no such floor, so that its bracket starts at 0: the floor would
+      // forbid moving a parameter by less than minAlpha |g|_inf.
+      const Scalar min_alpha
+          = resetB && _alpha < _ls_opts.alpha0 ? Scalar(0) : _ls_opts.minAlpha;
+      retCode = WolfeLineSearch(_func, _alpha, _xk_1, _fk_1, _gk_1, _pk, _xk,
+                                _fk, _gk, _ls_opts.c1, _ls_opts.c2, min_alpha,
+                                _ls_opts.maxLSIts, _ls_opts.maxLSRestarts);
       if (retCode) {
         // Line search failed...
         if (resetB) {
@@ -246,16 +264,25 @@ class BFGSMinimizer {
     // Compute search direction for next step
     _qn.search_direction(_pk, _gk);
 
+    // The first trial step was cut short by maxStep0. The tests on the change
+    // of f and x, and the relative gradient test, would judge convergence
+    // from that short step alone: a large gradient in one parameter makes
+    // the step in all others tiny.
+    const bool capped = resetB && _alpha0 < _ls_opts.alpha0;
+    _capped = capped && _itNum == 1;
+
     // Check for convergence
-    if (std::fabs(_fk_1 - _fk) < _conv_opts.tolAbsF) {
+    if (!capped && std::fabs(_fk_1 - _fk) < _conv_opts.tolAbsF) {
       // Objective function improvement wasn't sufficient
       retCode = TERM_ABSF;
     } else if (gradNorm < _conv_opts.tolAbsGrad) {
       retCode = TERM_ABSGRAD;  // Gradient norm was below threshold
-    } else if (stepNorm < _conv_opts.tolAbsX) {
+    } else if (!capped && stepNorm < _conv_opts.tolAbsX) {
       retCode = TERM_ABSX;  // Change in x was too small
     } else if (_itNum >= _conv_opts.maxIts) {
       retCode = TERM_MAXIT;  // Max number of iterations hit
+    } else if (capped) {
+      retCode = TERM_SUCCESS;
     } else if (rel_obj_decrease()
                < _conv_opts.tolRelF * std::numeric_limits<Scalar>::epsilon()) {
       // Relative improvement in objective function wasn't sufficient

src/stan/optimization/bfgs_linesearch.hpp

@@ -245,6 +245,11 @@ int WolfeLineSearch(FunctorType &func, Scalar &alpha, XType &x1,
   Scalar prevDFp(dfp);
   Scalar newDFp;
 
+  // the zoom gives up on a bracket narrower than 1e-16, or than 1e-13 of the
+  // first trial step when that is shorter than 1e-3 (a capped first step, or
+  // a short interpolated one)
+  const Scalar minRange(std::min(Scalar(1e-16), Scalar(1e-13) * alpha));
+
   int retCode = 0, nits = 0, lsRestarts = 0, ret;
 
   while (1) {
@@ -271,7 +276,7 @@ int WolfeLineSearch(FunctorType &func, Scalar &alpha, XType &x1,
     if ((func_val > f0 + alpha * c1dfp) || (func_val >= prevF && nits > 0)) {
       retCode = WolfLSZoom(alpha, x1, func_val, gradx1, func, x0, f0, dfp,
                            c1dfp, c2dfp, p, alpha0, prevF, prevDFp, alpha1,
-                           func_val, newDFp, 1e-16);
+                           func_val, newDFp, minRange);
       break;
     }
     if (std::fabs(newDFp) <= -c2dfp) {
@@ -281,7 +286,7 @@ int WolfeLineSearch(FunctorType &func, Scalar &alpha, XType &x1,
     if (newDFp >= 0) {
       retCode = WolfLSZoom(alpha, x1, func_val, gradx1, func, x0, f0, dfp,
                            c1dfp, c2dfp, p, alpha1, func_val, newDFp, alpha0,
-                           prevF, prevDFp, 1e-16);
+                           prevF, prevDFp, minRange);
       break;
     }
 

src/stan/services/optimize/bfgs.hpp

@@ -29,7 +29,8 @@ namespace optimize {
  * @param[in] random_seed random seed for the random number generator
  * @param[in] chain chain id to advance the pseudo random number generator
  * @param[in] init_radius radius to initialize
- * @param[in] init_alpha line search step size for first iteration
+ * @param[in] init_alpha line search step size for first iteration, reduced
+ *   if needed so that the first trial step moves no parameter by more than 1
  * @param[in] tol_obj convergence tolerance on absolute changes in
  *   objective function value
  * @param[in] tol_rel_obj convergence tolerance on relative changes
@@ -75,6 +76,7 @@ int bfgs(Model& model, const stan::io::var_context& init,
       Optimizer;
   Optimizer bfgs(model, cont_vector, disc_vector, &bfgs_ss);
   bfgs._ls_opts.alpha0 = init_alpha;
+  bfgs._ls_opts.maxStep0 = 1;
   bfgs._conv_opts.tolAbsF = tol_obj;
   bfgs._conv_opts.tolRelF = tol_rel_obj;
   bfgs._conv_opts.tolAbsGrad = tol_grad;

src/stan/services/optimize/lbfgs.hpp

@@ -30,7 +30,8 @@ namespace optimize {
  * @param[in] chain chain id to advance the pseudo random number generator
  * @param[in] init_radius radius to initialize
  * @param[in] history_size amount of history to keep for L-BFGS
- * @param[in] init_alpha line search step size for first iteration
+ * @param[in] init_alpha line search step size for first iteration, reduced
+ *   if needed so that the first trial step moves no parameter by more than 1
  * @param[in] tol_obj convergence tolerance on absolute changes in
  *   objective function value
  * @param[in] tol_rel_obj convergence tolerance on relative changes
@@ -79,6 +80,7 @@ int lbfgs(Model& model, const stan::io::var_context& init,
   Optimizer lbfgs(model, cont_vector, disc_vector, &lbfgs_ss);
   lbfgs.get_qnupdate().set_history_size(history_size);
   lbfgs._ls_opts.alpha0 = init_alpha;
+  lbfgs._ls_opts.maxStep0 = 1;
   lbfgs._conv_opts.tolAbsF = tol_obj;
   lbfgs._conv_opts.tolRelF = tol_rel_obj;
   lbfgs._conv_opts.tolAbsGrad = tol_grad;

src/test/test-models/good/optimization/two_scale.stan

@@ -0,0 +1,13 @@
+/**
+ * Two parameters with curvatures 1e12 and 1. When x1 starts close to its
+ * optimum, its gradient still sets the length of the first step, which then
+ * hardly moves x2. Optimum: (x1, x2) = (0, 5).
+ */
+parameters {
+  real x1;
+  real x2;
+}
+model {
+  x1 ~ normal(0, 1e-6);
+  x2 ~ normal(5, 1);
+}

src/test/unit/optimization/bfgs_linesearch_test.cpp

@@ -239,6 +239,34 @@ TEST(OptimizationBfgsLinesearch, wolfeLineSearch_jump_terminates) {
   EXPECT_LT(func.n_evals_, 1000);
 }
 
+TEST(OptimizationBfgsLinesearch, wolfeLineSearch_tiny_first_step) {
+  // phi(x) = -x + 1e17 x^2 / 2. The first trial step 2.35e-17 (one unit
+  // along a gradient of 4.3e16) overshoots the minimiser 1e-17. The bracket
+  // it gives, of width 2.35e-17, is narrower than the fixed width 1e-16 at
+  // which the zoom used to give up. minAlpha = 1e-20 lies below the first
+  // trial step.
+  using stan::optimization::WolfeLineSearch;
+  const double c1 = 1e-4, c2 = 0.9;
+  auto func = [](const Eigen::Matrix<double, Eigen::Dynamic, 1> &x, double &f,
+                 Eigen::Matrix<double, Eigen::Dynamic, 1> &g) {
+    f = -x[0] + 0.5e17 * x[0] * x[0];
+    g.resize(1);
+    g[0] = -1 + 1e17 * x[0];
+    return 0;
+  };
+  Eigen::Matrix<double, -1, 1> x0(1), x1, p(1), gradx0, gradx1;
+  double f0, f1;
+  x0 << 0.0;
+  func(x0, f0, gradx0);
+  p << 1.0;
+  double alpha = 2.35e-17;
+  int ret = WolfeLineSearch(func, alpha, x1, f1, gradx1, p, x0, f0, gradx0, c1,
+                            c2, 1e-20, 20.0, 10.0);
+  EXPECT_EQ(0, ret);
+  EXPECT_LE(f1, f0 + c1 * alpha * p.dot(gradx0));
+  EXPECT_LE(std::fabs(p.dot(gradx1)), c2 * std::fabs(p.dot(gradx0)));
+}
+
 TEST(OptimizationBfgsLinesearch, WolfLSZoom_adjacent_bracket) {
   // alo and ahi are adjacent doubles and every trial step rounds to ahi;
   // this looped forever (#3229)

src/test/unit/optimization/bfgs_minimizer_test.cpp

@@ -349,3 +349,97 @@ TEST_F(OptimizationBfgsMinimizer, minimize) {
 
   EXPECT_FLOAT_EQ(bfgs.minimize(cont_vector), 31);
 }
+
+namespace {
+// f(x) = sum cosh(x_i), minimised at 0; the evaluation fails where cosh
+// overflows
+struct minimizer_cosh {
+  int operator()(const Eigen::VectorXd &x, double &f, Eigen::VectorXd &g) {
+    f = x.array().cosh().sum();
+    g = x.array().sinh().matrix();
+    return std::isfinite(f) && g.allFinite() ? 0 : 1;
+  }
+};
+}  // namespace
+
+TEST(OptimizationBfgsMinimizerCosh, large_initial_gradient) {
+  // At x0 = (30, 30) the gradient is sinh(30) = 5.3e12 per coordinate. With
+  // the default maxStep0 = +inf the first trial step is init_alpha * |g| =
+  // 5.3e9, where cosh overflows, and after halving that step 10 times the
+  // first line search fails. With maxStep0 = 1 (as the optimize services set
+  // it) the first trial moves no coordinate by more than 1.
+  stan::optimization::BFGSMinimizer<minimizer_cosh,
+                                    stan::optimization::LBFGSUpdate<> >
+      lbfgs(minimizer_cosh{});
+  lbfgs._ls_opts.maxStep0 = 1;
+  Eigen::VectorXd x0(2);
+  x0 << 30, 30;
+  lbfgs.initialize(x0);
+  int ret = lbfgs.step();
+  EXPECT_EQ(stan::optimization::TERM_SUCCESS, ret);
+  while (ret == stan::optimization::TERM_SUCCESS) {
+    ret = lbfgs.step();
+  }
+  EXPECT_GT(ret, 0);
+  EXPECT_NEAR(2, lbfgs.curr_f(), 1e-8);
+}
+
+namespace {
+// f(x) = exp(30 x) - 30 x, minimised at 0
+struct minimizer_exp {
+  int operator()(const Eigen::VectorXd &x, double &f, Eigen::VectorXd &g) {
+    f = std::exp(30 * x[0]) - 30 * x[0];
+    g.resize(1);
+    g[0] = 30 * (std::exp(30 * x[0]) - 1);
+    return std::isfinite(f) && g.allFinite() ? 0 : 1;
+  }
+};
+}  // namespace
+
+TEST(OptimizationBfgsMinimizerExp, no_relative_convergence_after_reset) {
+  // From x0 = 2 (gradient 3e27) the first line search fails with the default
+  // maxStep0 = +inf. With the first trial step capped at 1 it accepts x = 1,
+  // where the gradient is 3e14 and the inverse Hessian estimate from that one
+  // step is 3e-28. The relative gradient test (with CmdStan's default
+  // tolerance) must not stop the run there.
+  stan::optimization::BFGSMinimizer<minimizer_exp,
+                                    stan::optimization::LBFGSUpdate<> >
+      lbfgs(minimizer_exp{});
+  lbfgs._ls_opts.maxStep0 = 1;
+  lbfgs._conv_opts.tolRelGrad = 1e7;
+  Eigen::VectorXd x0(1);
+  x0 << 2;
+  lbfgs.initialize(x0);
+  int ret = stan::optimization::TERM_SUCCESS;
+  while (ret == stan::optimization::TERM_SUCCESS) {
+    ret = lbfgs.step();
+  }
+  EXPECT_GT(ret, 0);
+  EXPECT_NEAR(1, lbfgs.curr_f(), 1e-8);
+}
+
+namespace {
+// f(x) = 1e13 x^2 / 2, minimised at 0
+struct minimizer_stiff_quadratic {
+  int operator()(const Eigen::VectorXd &x, double &f, Eigen::VectorXd &g) {
+    f = 0.5e13 * x.squaredNorm();
+    g = 1e13 * x;
+    return 0;
+  }
+};
+}  // namespace
+
+TEST(OptimizationBfgsMinimizerStiffQuadratic, capped_step_without_floor) {
+  // From x0 = 0.08 (gradient 8e11) the capped first trial step 1.25e-12
+  // moves x by 1, past the minimiser at step 1e-13. That step lies below
+  // minAlpha = 1e-12, the smallest step the line search may return, so the
+  // first search used to fail.
+  stan::optimization::BFGSMinimizer<minimizer_stiff_quadratic,
+                                    stan::optimization::LBFGSUpdate<> >
+      lbfgs(minimizer_stiff_quadratic{});
+  lbfgs._ls_opts.maxStep0 = 1;
+  Eigen::VectorXd x0(1);
+  x0 << 0.08;
+  lbfgs.initialize(x0);
+  EXPECT_GE(lbfgs.step(), 0);
+}

src/test/unit/services/optimize/bfgs_jacobian_test.cpp

@@ -1,6 +1,7 @@
 #include <stan/services/optimize/bfgs.hpp>
 #include <gtest/gtest.h>
 #include <stan/io/empty_var_context.hpp>
+#include <stan/io/json/json_data.hpp>
 #include <test/test-models/good/optimization/simple_jacobian.hpp>
 #include <test/unit/services/instrumented_callbacks.hpp>
 #include <stan/callbacks/stream_writer.hpp>
@@ -66,3 +67,15 @@ TEST_F(ServicesOptimize, withoutJacobian) {
   EXPECT_NEAR(3, parameter.states_.back()[2], 0.0001);
   EXPECT_GT(interrupt.call_count(), 1);
 }
+
+TEST_F(ServicesOptimize, bfgs_large_initial_gradient) {
+  // as lbfgs_large_initial_gradient in lbfgs_jacobian_test.cpp
+  std::stringstream init_json("{\"sigma\": 485165195.40979}");
+  stan::json::json_data init_context(init_json);
+  stan::test::unit::instrumented_interrupt interrupt;
+  int return_code = stan::services::optimize::bfgs<stan_model, false>(
+      model, init_context, 0, 1, 0, 0.001, 1e-12, 10000, 1e-8, 10000000, 1e-8,
+      2000, false, 0, interrupt, logger, init, parameter);
+  EXPECT_EQ(0, return_code);
+  EXPECT_NEAR(3, parameter.states_.back()[2], 0.0001);
+}

src/test/unit/services/optimize/lbfgs_jacobian_test.cpp

@@ -1,6 +1,7 @@
 #include <stan/services/optimize/lbfgs.hpp>
 #include <gtest/gtest.h>
 #include <stan/io/empty_var_context.hpp>
+#include <stan/io/json/json_data.hpp>
 #include <test/test-models/good/optimization/simple_jacobian.hpp>
 #include <test/unit/services/instrumented_callbacks.hpp>
 #include <stan/callbacks/stream_writer.hpp>
@@ -65,3 +66,19 @@ TEST_F(ServicesOptimize, without_jacobian) {
   EXPECT_EQ("sigma", parameter.names_[2]);
   EXPECT_NEAR(3, parameter.states_.back()[2], 0.0001);
 }
+
+TEST_F(ServicesOptimize, lbfgs_large_initial_gradient) {
+  // At sigma = exp(20) the gradient in log(sigma) is (sigma - 3) * sigma =
+  // 2.4e17. The first trial step was init_alpha times that gradient, the
+  // step that is needed is below minAlpha = 1e-12, and the first line search
+  // failed (error code 70). The first trial step now moves log(sigma) by at
+  // most 1.
+  std::stringstream init_json("{\"sigma\": 485165195.40979}");
+  stan::json::json_data init_context(init_json);
+  stan::test::unit::instrumented_interrupt interrupt;
+  int return_code = stan::services::optimize::lbfgs<stan_model, false>(
+      model, init_context, 0, 1, 0, 5, 0.001, 1e-12, 10000, 1e-8, 10000000,
+      1e-8, 2000, false, 0, interrupt, logger, init, parameter);
+  EXPECT_EQ(0, return_code);
+  EXPECT_NEAR(3, parameter.states_.back()[2], 0.0001);
+}