Newton optimizer: handle flat directions without producing NaN

src/stan/optimization/newton.hpp

@@ -1,9 +1,11 @@
 #ifndef STAN_OPTIMIZATION_NEWTON_HPP
 #define STAN_OPTIMIZATION_NEWTON_HPP
 
-#include <stan/model/grad_hess_log_prob.hpp>
-#include <stan/model/log_prob_grad.hpp>
+#include <stan/math/rev.hpp>
+#include <stan/math/rev/functor/finite_diff_hessian_auto.hpp>
 #include <stan/math/prim/fun/Eigen.hpp>
+#include <cmath>
+#include <limits>
 #include <vector>
 
 namespace stan {
@@ -12,40 +14,93 @@ namespace optimization {
 typedef Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> matrix_d;
 typedef Eigen::Matrix<double, Eigen::Dynamic, 1> vector_d;
 
-// Negates any positive eigenvalues in H so that H is negative
-// definite, and then solves Hu = g and stores the result into
-// g. Avoids problems due to non-log-concave distributions.
-inline void make_negative_definite_and_solve(matrix_d& H, vector_d& g) {
+/**
+ * Negates any positive eigenvalues in H so that H is negative
+ * definite, then solves Hu = g and stores the result into g.
+ * Avoids problems due to non-log-concave distributions.
+ *
+ * Each eigenvalue magnitude is floored at delta before inverting, so the
+ * step along a direction with little or no curvature is a gradient step
+ * scaled by 1 / delta rather than an unbounded or undefined quantity.
+ * This "saturating inverse" is continuous in the eigenvalues and bounds
+ * the effective condition number of the solve by 1 / sqrt(u).
+ *
+ * The floor is delta = max(sqrt(u) * max|lambda|, sqrt(u)), with u the
+ * unit roundoff. The relative term follows Nocedal and Wright, Numerical
+ * Optimization, 2nd ed., Section 3.4, which replaces problem eigenvalues
+ * with a delta of order sqrt(u). The absolute term is the same value
+ * under a well-scaled assumption and keeps an all-zero Hessian well
+ * defined. The backtracking line search in newton_step shortens any
+ * step that turns out too long.
+ *
+ * @param[in] H Hessian of the log density
+ * @param[in, out] g gradient on input, Newton step direction on output
+ */
+template <typename VecG>
+inline void make_negative_definite_and_solve(matrix_d& H, VecG& g) {
   Eigen::SelfAdjointEigenSolver<matrix_d> solver(H);
-  matrix_d eigenvectors = solver.eigenvectors();
+  auto&& eigenvectors = solver.eigenvectors();
   vector_d eigenvalues = solver.eigenvalues();
   vector_d eigenprojections = eigenvectors.transpose() * g;
+  const double sqrt_eps = std::sqrt(std::numeric_limits<double>::epsilon());
+  double max_abs_eigenvalue = eigenvalues.cwiseAbs().maxCoeff();
+  double delta = std::fmax(sqrt_eps * max_abs_eigenvalue, sqrt_eps);
   for (int i = 0; i < g.size(); i++) {
-    eigenprojections[i] = -eigenprojections[i] / fabs(eigenvalues[i]);
+    eigenprojections[i]
+        = -eigenprojections[i] / std::fmax(std::fabs(eigenvalues[i]), delta);
   }
   g = eigenvectors * eigenprojections;
 }
 
+/**
+ * Take one Newton step on the log density of the model, updating
+ * params_r in place.
+ *
+ * The gradient is computed by reverse-mode autodiff and the Hessian
+ * by central finite differences of the gradient with a per-coordinate
+ * step size, which costs 2 * params_r.size() + 1 gradient evaluations.
+ * The Hessian is made negative definite before solving for the Newton
+ * direction, and a backtracking line search on the log density chooses
+ * the step length. The log density is evaluated with all constant
+ * terms included (propto = false) so that the line search can use
+ * plain double arithmetic without autodiff.
+ *
+ * @tparam M Class of model.
+ * @tparam jacobian True if the log absolute Jacobian determinant of
+ * the inverse parameter transforms is added to the log density.
+ * @param[in] model Model.
+ * @param[in, out] params_r Unconstrained parameters; updated to the
+ * new point if the step improves the log density.
+ * @param[in] params_i Integer-valued parameters (unused).
+ * @param[in, out] output_stream Stream to which print statements in
+ * the Stan program are written.
+ * @return Log density, including constant terms, at the returned
+ * params_r.
+ */
 template <typename M, bool jacobian = false>
 double newton_step(M& model, std::vector<double>& params_r,
                    std::vector<int>& params_i,
                    std::ostream* output_stream = 0) {
-  std::vector<double> gradient;
-  std::vector<double> hessian;
+  const Eigen::Index n = params_r.size();
+  const vector_d x = Eigen::Map<const vector_d>(params_r.data(), n);
 
-  double f0 = stan::model::grad_hess_log_prob<true, jacobian>(
-      model, params_r, params_i, gradient, hessian);
-  matrix_d H(params_r.size(), params_r.size());
-  for (size_t i = 0; i < hessian.size(); i++) {
-    H(i) = hessian[i];
+  auto log_density = [&](auto&& theta) {
+    return model.template log_prob<false, jacobian, stan::math::var>(
+        theta, output_stream);
+  };
+  double f0;
+  vector_d g;
+  matrix_d H;
+  stan::math::internal::finite_diff_hessian_auto(log_density, x, f0, g, H);
+  if (!std::isfinite(f0)) {
+    return f0;
   }
-  vector_d g(params_r.size());
-  for (size_t i = 0; i < gradient.size(); i++)
-    g(i) = gradient[i];
   make_negative_definite_and_solve(H, g);
-  //         H.ldlt().solveInPlace(g);
+  if (!g.array().allFinite()) {
+    return f0;
+  }
 
-  std::vector<double> new_params_r(params_r.size());
+  vector_d new_params_r(n);
   double step_size = 2;
   double min_step_size = 1e-50;
   double f1 = -1e100;
@@ -55,18 +110,23 @@ double newton_step(M& model, std::vector<double>& params_r,
     if (step_size < min_step_size)
       return f0;
 
-    for (size_t i = 0; i < params_r.size(); i++)
-      new_params_r[i] = params_r[i] - step_size * g[i];
+    new_params_r = x - step_size * g;
+    if (!new_params_r.array().allFinite()) {
+      f1 = -1e100;
+      continue;
+    }
     try {
-      f1 = stan::model::log_prob_grad<true, jacobian>(model, new_params_r,
-                                                      params_i, gradient);
+      f1 = model.template log_prob<false, jacobian, double>(new_params_r,
+                                                            output_stream);
     } catch (std::domain_error& e) {
       // FIXME:  this is not a good way to handle a general exception
       f1 = -1e100;
     }
+    if (!std::isfinite(f1)) {
+      f1 = -1e100;
+    }
   }
-  for (size_t i = 0; i < params_r.size(); i++)
-    params_r[i] = new_params_r[i];
+  Eigen::Map<vector_d>(params_r.data(), n) = new_params_r;
 
   return f1;
 }

src/stan/services/optimize/newton.hpp

@@ -36,7 +36,8 @@ namespace optimize {
  * @param[in,out] logger Logger for messages
  * @param[in,out] init_writer Writer callback for unconstrained inits
  * @param[in,out] parameter_writer output for parameter values
- * @return error_codes::OK if successful
+ * @return error_codes::OK if successful, error_codes::SOFTWARE if the
+ *   final log probability or parameters are not finite
  */
 template <class Model, bool jacobian = false>
 int newton(Model& model, const stan::io::var_context& init,
@@ -120,7 +121,15 @@ int newton(Model& model, const stan::io::var_context& init,
       break;
   }
 
-  if (std::fabs(lp - lastlp) <= 1e-8) {
+  bool finite_result
+      = std::isfinite(lp)
+        && Eigen::Map<const vector_d>(cont_vector.data(), cont_vector.size())
+               .array()
+               .isFinite()
+               .all();
+  if (!finite_result) {
+    ret = optimization::TERM_LSFAIL;
+  } else if (std::fabs(lp - lastlp) <= 1e-8) {
     ret = optimization::TERM_ABSF;
   } else {
     ret = optimization::TERM_MAXIT;
@@ -135,6 +144,13 @@ int newton(Model& model, const stan::io::var_context& init,
     values.insert(values.begin(), {lp, static_cast<double>(ret)});
     parameter_writer(values);
   }
+
+  if (!finite_result) {
+    logger.error(
+        "Optimization terminated with error: "
+        "log probability or parameters are not finite.");
+    return error_codes::SOFTWARE;
+  }
   return error_codes::OK;
 }
 

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

@@ -0,0 +1,12 @@
+/**
+ * The target does not depend on x, so the gradient and Hessian
+ * are identically zero along that direction. Used to check that
+ * the Newton optimizer handles a flat direction without producing
+ * non-finite parameter values.
+ */
+parameters {
+  real x;
+}
+model {
+  target += 0.5;
+}

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

@@ -0,0 +1,11 @@
+/**
+ * The target is linear in x, so the gradient is nonzero while the
+ * Hessian is identically zero. Used to check that the Newton optimizer
+ * still moves along a direction with no curvature.
+ */
+parameters {
+  real x;
+}
+model {
+  target += x;
+}

src/test/unit/optimization/newton_linear_target_test.cpp

@@ -0,0 +1,28 @@
+#include <gtest/gtest.h>
+#include <stan/optimization/newton.hpp>
+#include <stan/io/empty_var_context.hpp>
+#include <test/test-models/good/optimization/linear_target.hpp>
+#include <cmath>
+#include <limits>
+#include <vector>
+
+typedef linear_target_model_namespace::linear_target_model Model;
+
+TEST(OptimizationNewton, linear_target_moves_uphill_by_bounded_step) {
+  const double sqrt_eps = std::sqrt(std::numeric_limits<double>::epsilon());
+  stan::io::empty_var_context dummy_context;
+  Model model(dummy_context);
+
+  std::vector<double> params_r(1, 0.0);
+  std::vector<int> params_i;
+
+  double f = stan::optimization::newton_step<Model, false>(model, params_r,
+                                                           params_i);
+
+  ASSERT_EQ(1u, params_r.size());
+  EXPECT_TRUE(std::isfinite(params_r[0]));
+  EXPECT_GT(params_r[0], 0.0) << "zero-curvature direction must still move";
+  EXPECT_LE(params_r[0], 2.0 / sqrt_eps)
+      << "step along a zero-curvature direction must be bounded by the floor";
+  EXPECT_GT(f, 0.0) << "objective must improve along a linear target";
+}

src/test/unit/optimization/newton_test.cpp

@@ -0,0 +1,101 @@
+#include <gtest/gtest.h>
+#include <stan/optimization/newton.hpp>
+#include <stan/io/empty_var_context.hpp>
+#include <test/test-models/good/optimization/flat_target.hpp>
+#include <cmath>
+#include <limits>
+#include <vector>
+
+typedef flat_target_model_namespace::flat_target_model Model;
+
+// Regression test for https://github.com/stan-dev/stan/issues/3425
+TEST(OptimizationNewton, flat_direction_keeps_parameters_finite) {
+  stan::io::empty_var_context dummy_context;
+  Model model(dummy_context);
+
+  std::vector<double> params_r(1, 1.0);
+  std::vector<int> params_i;
+
+  double f = stan::optimization::newton_step<Model, false>(model, params_r,
+                                                           params_i);
+
+  EXPECT_FLOAT_EQ(0.5, f);
+  ASSERT_EQ(1u, params_r.size());
+  EXPECT_TRUE(std::isfinite(params_r[0]))
+      << "newton_step produced non-finite parameter: " << params_r[0];
+}
+
+TEST(OptimizationNewton,
+     make_negative_definite_and_solve_floors_small_eigenvalue_at_sqrt_eps) {
+  const double eps = std::numeric_limits<double>::epsilon();
+  const double sqrt_eps = std::sqrt(eps);
+  stan::optimization::matrix_d H = stan::optimization::matrix_d::Zero(2, 2);
+  H(0, 0) = -1.0;
+  H(1, 1) = -4.0 * eps;
+  stan::optimization::vector_d g = stan::optimization::vector_d::Ones(2);
+
+  stan::optimization::make_negative_definite_and_solve(H, g);
+
+  EXPECT_FLOAT_EQ(-1.0, g[0]);
+  EXPECT_FLOAT_EQ(-1.0 / sqrt_eps, g[1])
+      << "eigenvalue below sqrt(eps) * max should be floored, not dropped";
+}
+
+TEST(OptimizationNewton,
+     make_negative_definite_and_solve_zero_eigenvalue_uses_relative_floor) {
+  const double sqrt_eps = std::sqrt(std::numeric_limits<double>::epsilon());
+  stan::optimization::matrix_d H = stan::optimization::matrix_d::Zero(2, 2);
+  H(0, 0) = -1.0;
+  stan::optimization::vector_d g = stan::optimization::vector_d::Ones(2);
+
+  stan::optimization::make_negative_definite_and_solve(H, g);
+
+  EXPECT_FLOAT_EQ(-1.0, g[0]);
+  EXPECT_FLOAT_EQ(-1.0 / sqrt_eps, g[1]);
+}
+
+TEST(OptimizationNewton,
+     make_negative_definite_and_solve_is_continuous_at_old_cutoff) {
+  const double eps = std::numeric_limits<double>::epsilon();
+  const double old_cutoff = 4.0 * 2 * eps;
+  stan::optimization::matrix_d H_above
+      = stan::optimization::matrix_d::Zero(2, 2);
+  H_above(0, 0) = -1.0;
+  H_above(1, 1) = -1.01 * old_cutoff;
+  stan::optimization::matrix_d H_below = H_above;
+  H_below(1, 1) = -0.99 * old_cutoff;
+  stan::optimization::vector_d g_above = stan::optimization::vector_d::Ones(2);
+  stan::optimization::vector_d g_below = g_above;
+
+  stan::optimization::make_negative_definite_and_solve(H_above, g_above);
+  stan::optimization::make_negative_definite_and_solve(H_below, g_below);
+
+  EXPECT_NEAR(g_above[1], g_below[1], 1e-6 * std::fabs(g_above[1]))
+      << "step must not jump when an eigenvalue crosses the cutoff";
+}
+
+TEST(OptimizationNewton,
+     make_negative_definite_and_solve_zero_hessian_nonzero_gradient) {
+  const double sqrt_eps = std::sqrt(std::numeric_limits<double>::epsilon());
+  stan::optimization::matrix_d H = stan::optimization::matrix_d::Zero(2, 2);
+  stan::optimization::vector_d g = stan::optimization::vector_d::Ones(2);
+
+  stan::optimization::make_negative_definite_and_solve(H, g);
+
+  for (int i = 0; i < g.size(); ++i) {
+    EXPECT_FLOAT_EQ(-1.0 / sqrt_eps, g[i])
+        << "all-zero Hessian must use the absolute floor, component " << i;
+  }
+}
+
+TEST(OptimizationNewton, make_negative_definite_and_solve_zero_hessian) {
+  stan::optimization::matrix_d H = stan::optimization::matrix_d::Zero(2, 2);
+  stan::optimization::vector_d g = stan::optimization::vector_d::Zero(2);
+
+  stan::optimization::make_negative_definite_and_solve(H, g);
+
+  for (int i = 0; i < g.size(); ++i) {
+    EXPECT_TRUE(std::isfinite(g[i]))
+        << "step direction has non-finite component " << i << ": " << g[i];
+  }
+}

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

@@ -0,0 +1,44 @@
+#include <stan/services/optimize/newton.hpp>
+#include <gtest/gtest.h>
+#include <stan/io/empty_var_context.hpp>
+#include <test/test-models/good/optimization/flat_target.hpp>
+#include <test/unit/services/instrumented_callbacks.hpp>
+#include <stan/callbacks/stream_writer.hpp>
+#include <cmath>
+
+struct ServicesOptimizeNewtonFlatTarget : public testing::Test {
+  ServicesOptimizeNewtonFlatTarget()
+      : init(init_ss), parameter(parameter_ss), model(context, 0, &model_ss) {}
+
+  std::stringstream init_ss, parameter_ss, model_ss;
+  stan::test::unit::instrumented_logger logger;
+  stan::callbacks::stream_writer init;
+  stan::test::unit::values_writer parameter;
+  stan::io::empty_var_context context;
+  stan_model model;
+};
+
+// Regression test for https://github.com/stan-dev/stan/issues/3425
+// The service must not report success while writing non-finite parameters.
+TEST_F(ServicesOptimizeNewtonFlatTarget, does_not_report_ok_with_nan_params) {
+  unsigned int seed = 0;
+  unsigned int chain = 1;
+  double init_radius = 1;
+  int num_iterations = 10;
+  bool save_iterations = false;
+  stan::test::unit::instrumented_interrupt interrupt;
+
+  int return_code = stan::services::optimize::newton(
+      model, context, seed, chain, init_radius, num_iterations, save_iterations,
+      interrupt, logger, init, parameter);
+
+  ASSERT_EQ(3, parameter.names_.size());
+  EXPECT_EQ("x", parameter.names_[2]);
+  ASSERT_EQ(1, parameter.states_.size());
+
+  double x = parameter.states_.back()[2];
+  EXPECT_TRUE(std::isfinite(x)
+              || return_code != stan::services::error_codes::OK)
+      << "newton returned error_codes::OK with x = " << x;
+  EXPECT_TRUE(std::isfinite(x)) << "final x = " << x;
+}