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,55 @@ 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.
+ *
+ * Eigenvalues whose magnitude is negligible relative to the largest
+ * eigenvalue are treated as zero and their directions are dropped
+ * from the solve, as in a pseudo-inverse. This keeps the step finite
+ * when the target is flat along some direction.
+ *
+ * @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;
+  double max_abs_eigenvalue = eigenvalues.cwiseAbs().maxCoeff();
+  double tolerance
+      = max_abs_eigenvalue * H.rows() * std::numeric_limits<double>::epsilon();
   for (int i = 0; i < g.size(); i++) {
-    eigenprojections[i] = -eigenprojections[i] / fabs(eigenvalues[i]);
+    double abs_eigenvalue = std::fabs(eigenvalues[i]);
+    if (abs_eigenvalue <= tolerance) {
+      eigenprojections[i] = 0;
+    } else {
+      eigenprojections[i] = -eigenprojections[i] / abs_eigenvalue;
+    }
   }
   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 +72,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 +83,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 +99,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/unit/optimization/newton_test.cpp

@@ -0,0 +1,37 @@
+#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 <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_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;
+}