Newton optimizer: handle flat directions without producing NaN

src/stan/optimization/newton.hpp

@@ -4,6 +4,8 @@
 #include <stan/model/grad_hess_log_prob.hpp>
 #include <stan/model/log_prob_grad.hpp>
 #include <stan/math/prim/fun/Eigen.hpp>
+#include <cmath>
+#include <limits>
 #include <vector>
 
 namespace stan {
@@ -12,20 +14,60 @@ 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.
+/**
+ * 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
+ */
 inline void make_negative_definite_and_solve(matrix_d& H, vector_d& g) {
   Eigen::SelfAdjointEigenSolver<matrix_d> solver(H);
   matrix_d 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;
 }
 
+/**
+ * Returns true if every element of the vector is finite.
+ *
+ * @tparam Vec vector type with size() and operator[]
+ * @param[in] v vector to check
+ * @return true if all elements are finite
+ */
+template <typename Vec>
+inline bool all_finite(const Vec& v) {
+  for (int i = 0; i < static_cast<int>(v.size()); ++i) {
+    if (!std::isfinite(v[i])) {
+      return false;
+    }
+  }
+  return true;
+}
+
 template <typename M, bool jacobian = false>
 double newton_step(M& model, std::vector<double>& params_r,
                    std::vector<int>& params_i,
@@ -35,6 +77,9 @@ double newton_step(M& model, std::vector<double>& params_r,
 
   double f0 = stan::model::grad_hess_log_prob<true, jacobian>(
       model, params_r, params_i, gradient, hessian);
+  if (!std::isfinite(f0)) {
+    return f0;
+  }
   matrix_d H(params_r.size(), params_r.size());
   for (size_t i = 0; i < hessian.size(); i++) {
     H(i) = hessian[i];
@@ -43,7 +88,9 @@ double newton_step(M& model, std::vector<double>& params_r,
   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 (!all_finite(g)) {
+    return f0;
+  }
 
   std::vector<double> new_params_r(params_r.size());
   double step_size = 2;
@@ -57,13 +104,20 @@ double newton_step(M& model, std::vector<double>& params_r,
 
     for (size_t i = 0; i < params_r.size(); i++)
       new_params_r[i] = params_r[i] - step_size * g[i];
+    if (!all_finite(new_params_r)) {
+      f1 = -1e100;
+      continue;
+    }
     try {
       f1 = stan::model::log_prob_grad<true, jacobian>(model, new_params_r,
                                                       params_i, gradient);
     } 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];

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,11 @@ int newton(Model& model, const stan::io::var_context& init,
       break;
   }
 
-  if (std::fabs(lp - lastlp) <= 1e-8) {
+  bool finite_result
+      = std::isfinite(lp) && optimization::all_finite(cont_vector);
+  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 +140,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;
+}