Fix NaN in hmm_hidden_state_prob for long sequences (#2677)

stan/math/prim/err/hmm_check.hpp

@@ -19,8 +19,9 @@ namespace math {
  * @param[in] Gamma transition density between hidden states.
  * @param[in] rho initial state
  * @param[in] function the name of the function using the arguments.
- * @throw `std::invalid_argument` if Gamma is not square
- *         or if the size of rho is not the number of rows of log_omegas.
+ * @throw `std::invalid_argument` if Gamma is not square, if log_omegas
+ *         has size 0, or if the size of rho is not the number of rows of
+ *         log_omegas.
  * @throw `std::domain_error` if rho is not a simplex or the rows
  *         of Gamma are not a simplex.
  */
@@ -34,6 +35,7 @@ inline void hmm_check(const T_omega& log_omegas, const T_Gamma& Gamma,
   check_consistent_size(function, "rho", rho, n_states);
   check_square(function, "Gamma", Gamma);
   check_nonzero_size(function, "Gamma", Gamma);
+  check_nonzero_size(function, "log_omegas", log_omegas);
   check_multiplicable(function, "Gamma", Gamma, "log_omegas", log_omegas);
 
   check_simplex(function, "rho", rho);

stan/math/prim/prob/hmm_hidden_state_prob.hpp

@@ -20,15 +20,21 @@ namespace math {
  * Hence for each hidden variable x, we compute a simplex with K elements.
  * The final result is stored in a K by N matrix, where N is the length of x.
  * log_omegas is a matrix of observational densities, where
- * the (i, j)th entry corresponds to the density of the ith observation, y_i,
- * given x_i = j.
+ * the (i, j)th entry corresponds to the density of the jth observation, y_j,
+ * given x_j = i.
  * The transition matrix Gamma is such that the (i, j)th entry is the
  * probability that x_n = j given x_{n - 1} = i. The rows of Gamma are
  * simplexes.
  * This function cannot be used to reconstruct the marginal distributon
  * of a state sequence given parameters and an observation sequence,
  * p(x | y, theta),
  * because it only computes marginals on a state-by-state basis.
+ * Each column of log_omegas is shifted by its maximum before it is
+ * exponentiated, and the forward and backward passes are normalized at
+ * every step, so long sequences and log densities of large magnitude do not
+ * underflow or overflow. The result is still NaN if the largest density in
+ * a column belongs to a state that the chain cannot reach, and the log
+ * densities of the reachable states are more than about 745 below it.
  *
  * @tparam T_omega type of the log likelihood matrix
  * @tparam T_Gamma type of the transition matrix
@@ -51,18 +57,30 @@ inline Eigen::MatrixXd hmm_hidden_state_prob(const T_omega& log_omegas,
   int n_states = log_omegas.rows();
   int n_transitions = log_omegas.cols() - 1;
 
-  Eigen::MatrixXd omegas = value_of(log_omegas).array().exp();
   ref_type_t<decltype(value_of(rho))> rho_dbl = value_of(rho);
   ref_type_t<decltype(value_of(Gamma))> Gamma_dbl = value_of(Gamma);
   hmm_check(log_omegas, Gamma_dbl, rho_dbl, "hmm_hidden_state_prob");
 
+  // Subtract the maximum of each column before exponentiating, so that
+  // omegas neither underflows nor overflows. The hidden state probabilities
+  // are invariant to a positive factor on each column of omegas.
+  ref_type_t<decltype(value_of(log_omegas))> log_omegas_dbl
+      = value_of(log_omegas);
+  Eigen::MatrixXd omegas
+      = (log_omegas_dbl.rowwise() - log_omegas_dbl.colwise().maxCoeff())
+            .array()
+            .exp();
+
   Eigen::MatrixXd alphas(n_states, n_transitions + 1);
   alphas.col(0) = omegas.col(0).cwiseProduct(rho_dbl);
   alphas.col(0) /= alphas.col(0).maxCoeff();
 
-  for (int n = 0; n < n_transitions; ++n)
+  // Forward pass with running normalization
+  for (int n = 0; n < n_transitions; ++n) {
     alphas.col(n + 1)
         = omegas.col(n + 1).cwiseProduct(Gamma_dbl.transpose() * alphas.col(n));
+    alphas.col(n + 1) /= alphas.col(n + 1).maxCoeff();
+  }
 
   // Backward pass with running normalization
   Eigen::VectorXd beta = Eigen::VectorXd::Ones(n_states);

stan/math/prim/prob/hmm_latent_rng.hpp

@@ -18,8 +18,8 @@ namespace math {
  * In this setting, the hidden states are discrete
  * and takes values over the finite space {1, ..., K}.
  * log_omegas is a matrix of observational densities, where
- * the (i, j)th entry corresponds to the density of the ith observation, y_i,
- * given x_i = j.
+ * the (i, j)th entry corresponds to the density of the jth observation, y_j,
+ * given x_j = i.
  * The transition matrix Gamma is such that the (i, j)th entry is the
  * probability that x_n = j given x_{n - 1} = i. The rows of Gamma are
  * simplexes.
@@ -47,11 +47,20 @@ inline std::vector<int> hmm_latent_rng(const T_omega& log_omegas,
   int n_states = log_omegas.rows();
   int n_transitions = log_omegas.cols() - 1;
 
-  Eigen::MatrixXd omegas = value_of(log_omegas).array().exp();
   ref_type_t<decltype(value_of(rho))> rho_dbl = value_of(rho);
   ref_type_t<decltype(value_of(Gamma))> Gamma_dbl = value_of(Gamma);
   hmm_check(log_omegas, Gamma_dbl, rho_dbl, "hmm_latent_rng");
 
+  // Subtract the maximum of each column before exponentiating, so that
+  // omegas neither underflows nor overflows. The sampling probabilities
+  // are invariant to a positive factor on each column of omegas.
+  ref_type_t<decltype(value_of(log_omegas))> log_omegas_dbl
+      = value_of(log_omegas);
+  Eigen::MatrixXd omegas
+      = (log_omegas_dbl.rowwise() - log_omegas_dbl.colwise().maxCoeff())
+            .array()
+            .exp();
+
   Eigen::MatrixXd alphas(n_states, n_transitions + 1);
   alphas.col(0) = omegas.col(0).cwiseProduct(rho_dbl);
   alphas.col(0) /= alphas.col(0).maxCoeff();

stan/math/prim/prob/hmm_marginal.hpp

@@ -98,8 +98,15 @@ inline auto hmm_marginal(const T_omega& log_omegas, const T_Gamma& Gamma,
 
   eig_matrix_partial alphas(n_states, n_transitions + 1);
   eig_vector_partial alpha_log_norms(n_transitions + 1);
+  // Subtract the maximum of each column before exponentiating, so that
+  // omegas neither underflows nor overflows. The shifts are added back to
+  // the log density. The gradients are invariant to the shifts.
+  const auto& log_omegas_val = to_ref(value_of(log_omegas_ref));
+  Eigen::Matrix<T_partial_type, 1, Eigen::Dynamic> log_omegas_max
+      = log_omegas_val.colwise().maxCoeff();
+  eig_matrix_partial omegas
+      = (log_omegas_val.rowwise() - log_omegas_max).array().exp();
   // compute the density using the forward algorithm.
-  eig_matrix_partial omegas = value_of(log_omegas_ref).array().exp();
   T_partial_type norm_norm;
   auto log_marginal_density = hmm_marginal_val(
       omegas, Gamma_val, rho_val, alphas, alpha_log_norms, norm_norm);
@@ -148,7 +155,7 @@ inline auto hmm_marginal(const T_omega& log_omegas, const T_Gamma& Gamma,
         edge<2>(ops_partials).partials_
             = omegas.col(0) / exp(log_marginal_density);
       }
-      return ops_partials.build(log_marginal_density);
+      return ops_partials.build(log_marginal_density + log_omegas_max.sum());
     } else {
       auto grad_corr_boundary = exp(kappa_log_norms(0) - norm_norm);
       eig_vector_partial C = Gamma_val * omegas.col(1).cwiseProduct(kappa[0]);
@@ -176,7 +183,7 @@ inline auto hmm_marginal(const T_omega& log_omegas, const T_Gamma& Gamma,
     }
   }
 
-  return ops_partials.build(log_marginal_density);
+  return ops_partials.build(log_marginal_density + log_omegas_max.sum());
 }
 
 }  // namespace math

test/unit/math/prim/err/hmm_check_test.cpp

@@ -53,6 +53,13 @@ TEST(err, hmm_check) {
       "hmm_marginal_lpdf: Columns of Gamma (3)"
       " and Rows of log_omegas (2) must match in size")
 
+  // log_omegas has no columns, so there are no observations.
+  MatrixXd log_omegas_empty(n_states, 0);
+  EXPECT_THROW_MSG(
+      hmm_check(log_omegas_empty, Gamma, rho, "hmm_marginal_lpdf"),
+      std::invalid_argument,
+      "hmm_marginal_lpdf: log_omegas has size 0, but must have a non-zero size")
+
   // rho is not a simplex.
   VectorXd rho_bad = rho;
   rho_bad(0) = rho(0) + 1;

test/unit/math/prim/prob/hmm_hidden_state_prob_test.cpp

@@ -1,13 +1,56 @@
 #include <test/unit/math/prim/prob/hmm_util.hpp>
 #include <stan/math/prim/prob/hmm_hidden_state_prob.hpp>
+#include <stan/math/prim/fun/log_sum_exp.hpp>
 #include <boost/math/distributions.hpp>
 #include <boost/random.hpp>
 #include <test/unit/math/test_ad.hpp>
 #include <test/unit/util.hpp>
 #include <gtest/gtest.h>
+#include <cmath>
 #include <limits>
 #include <vector>
 
+namespace {
+/**
+ * Reference implementation of hmm_hidden_state_prob: the forward-backward
+ * algorithm computed entirely on the log scale, with log_sum_exp in place
+ * of every sum, so that neither the forward nor the backward pass can
+ * underflow or overflow.
+ */
+Eigen::MatrixXd hmm_hidden_state_prob_log_space(
+    const Eigen::MatrixXd& log_omegas, const Eigen::MatrixXd& Gamma,
+    const Eigen::VectorXd& rho) {
+  using stan::math::log_sum_exp;
+  const int n_states = log_omegas.rows();
+  const int n_cols = log_omegas.cols();
+  const Eigen::MatrixXd log_Gamma = Gamma.array().log();
+
+  Eigen::MatrixXd log_alphas(n_states, n_cols);
+  log_alphas.col(0) = rho.array().log() + log_omegas.col(0).array();
+  for (int n = 1; n < n_cols; ++n) {
+    for (int j = 0; j < n_states; ++j) {
+      Eigen::VectorXd terms = log_alphas.col(n - 1) + log_Gamma.col(j);
+      log_alphas(j, n) = log_omegas(j, n) + log_sum_exp(terms);
+    }
+  }
+
+  Eigen::MatrixXd prob(n_states, n_cols);
+  Eigen::VectorXd log_beta = Eigen::VectorXd::Zero(n_states);
+  for (int n = n_cols - 1; n >= 0; --n) {
+    if (n < n_cols - 1) {
+      Eigen::VectorXd next = log_omegas.col(n + 1) + log_beta;
+      for (int i = 0; i < n_states; ++i) {
+        Eigen::VectorXd terms = log_Gamma.row(i).transpose() + next;
+        log_beta(i) = log_sum_exp(terms);
+      }
+    }
+    Eigen::VectorXd log_prob = log_alphas.col(n) + log_beta;
+    prob.col(n) = (log_prob.array() - log_sum_exp(log_prob)).exp();
+  }
+  return prob;
+}
+}  // namespace
+
 TEST_F(hmm_test, hidden_state_single_outcome) {
   using stan::math::hmm_hidden_state_prob;
 
@@ -19,7 +62,7 @@ TEST_F(hmm_test, hidden_state_single_outcome) {
 
   Eigen::MatrixXd prob = hmm_hidden_state_prob(log_omegas_, Gamma, rho);
 
-  for (int i = 0; i < n_transitions_; i++) {
+  for (int i = 0; i <= n_transitions_; i++) {
     EXPECT_EQ(prob(0, i), 1);
     EXPECT_EQ(prob(1, i), 0);
   }
@@ -37,7 +80,7 @@ TEST_F(hmm_test, hidden_state_identity_transition) {
 
   Eigen::MatrixXd prob = hmm_hidden_state_prob(log_omegas, Gamma, rho_);
 
-  for (int i = 0; i < n_transitions_; i++) {
+  for (int i = 0; i <= n_transitions_; i++) {
     EXPECT_FLOAT_EQ(prob(0, i), rho_(0));
     EXPECT_FLOAT_EQ(prob(1, i), rho_(1));
   }
@@ -64,8 +107,142 @@ TEST(hmm_test_nonstandard, hidden_state_symmetry) {
   EXPECT_FLOAT_EQ(prob(0, 0), 0.3);
   EXPECT_FLOAT_EQ(prob(1, 0), 0.7);
 
-  for (int i = 1; i < n_transitions; i++) {
+  for (int i = 1; i <= n_transitions; i++) {
     EXPECT_FLOAT_EQ(prob(0, i), 0.5);
     EXPECT_FLOAT_EQ(prob(1, i), 0.5);
   }
 }
+
+TEST(hmm_test_nonstandard, hidden_state_prob1) {
+  // This time, the transition matrix forces states to transition
+  // to the first state with probability 1. The large log densities
+  // overflow if they are exponentiated directly (#2677).
+  using stan::math::hmm_hidden_state_prob;
+  int n_states = 2;
+  int n_transitions = 2;
+  Eigen::MatrixXd Gamma(n_states, n_states);
+  Gamma << 1, 0, 1, 0;
+  Eigen::VectorXd rho(n_states);
+  rho << 0.3, 0.7;
+  Eigen::MatrixXd log_omegas
+      = 1000 * Eigen::MatrixXd::Ones(n_states, n_transitions + 1);
+
+  Eigen::MatrixXd prob = hmm_hidden_state_prob(log_omegas, Gamma, rho);
+  EXPECT_FLOAT_EQ(prob(0, 0), 0.3);
+  EXPECT_FLOAT_EQ(prob(1, 0), 0.7);
+
+  for (int i = 1; i <= n_transitions; i++) {
+    EXPECT_FLOAT_EQ(prob(0, i), 1);
+    EXPECT_FLOAT_EQ(prob(1, i), 0);
+  }
+}
+
+TEST_F(hmm_test, hidden_state_very_negative_log_omegas) {
+  // Regression test for #2677. Every log density is below -1000, so
+  // exponentiating it directly underflows to 0. Adding a constant to a
+  // column of log_omegas does not change the hidden state probabilities.
+  using stan::math::hmm_hidden_state_prob;
+  Eigen::MatrixXd log_omegas = log_omegas_;
+  for (int n = 0; n < log_omegas.cols(); ++n) {
+    log_omegas.col(n).array() -= 1000 + 10 * n;
+  }
+
+  Eigen::MatrixXd prob = hmm_hidden_state_prob(log_omegas, Gamma_, rho_);
+  Eigen::MatrixXd prob_unshifted
+      = hmm_hidden_state_prob(log_omegas_, Gamma_, rho_);
+  Eigen::MatrixXd prob_ref
+      = hmm_hidden_state_prob_log_space(log_omegas, Gamma_, rho_);
+
+  for (int n = 0; n < prob.cols(); ++n) {
+    for (int k = 0; k < prob.rows(); ++k) {
+      EXPECT_TRUE(std::isfinite(prob(k, n)));
+      EXPECT_NEAR(prob(k, n), prob_unshifted(k, n), 1e-12);
+      // The log scale reference loses about machine epsilon times the
+      // magnitude of the log densities, hence the looser tolerance.
+      EXPECT_NEAR(prob(k, n), prob_ref(k, n), 1e-10);
+    }
+  }
+}
+
+TEST(hmm_test_nonstandard, hidden_state_long_sequence) {
+  // Regression test for #2677. The log density of a long sequence is far
+  // below the smallest double, so the forward pass underflows to 0 unless
+  // it is rescaled at every step.
+  using stan::math::hmm_hidden_state_prob;
+  using stan::math::hmm_marginal;
+  int n_states = 2;
+  int n_obs = 5000;
+  Eigen::MatrixXd Gamma(n_states, n_states);
+  Gamma << 0.95, 0.05, 0.10, 0.90;
+  Eigen::VectorXd rho(n_states);
+  rho << 2.0 / 3, 1.0 / 3;
+
+  // Simulate a two-state Gaussian hidden Markov model, with observations
+  // normal(1, 1) in state 0 and normal(-1, 1) in state 1.
+  boost::random::mt19937 rng(2677);
+  boost::random::uniform_01<> uniform;
+  boost::random::normal_distribution<> normal(0, 1);
+  Eigen::MatrixXd log_omegas(n_states, n_obs);
+  int state = 0;
+  for (int n = 0; n < n_obs; ++n) {
+    state = uniform(rng) < Gamma(state, 0) ? 0 : 1;
+    double y = (state == 0 ? 1 : -1) + normal(rng);
+    for (int k = 0; k < n_states; ++k) {
+      log_omegas(k, n) = state_lpdf(y, 1, 1, k);
+    }
+  }
+  // The joint density underflows the double range many times over.
+  EXPECT_LT(hmm_marginal(log_omegas, Gamma, rho), -2000);
+
+  Eigen::MatrixXd prob = hmm_hidden_state_prob(log_omegas, Gamma, rho);
+  Eigen::MatrixXd prob_ref
+      = hmm_hidden_state_prob_log_space(log_omegas, Gamma, rho);
+
+  int first_non_finite = -1;
+  for (int n = 0; n < n_obs && first_non_finite < 0; ++n) {
+    if (!prob.col(n).allFinite()) {
+      first_non_finite = n;
+    }
+  }
+  ASSERT_EQ(first_non_finite, -1) << "first non-finite column";
+  EXPECT_LT((prob.colwise().sum().array() - 1).abs().maxCoeff(), 1e-12);
+  // The log scale reference loses about machine epsilon times the
+  // magnitude of the log densities, hence the looser tolerance.
+  EXPECT_LT((prob - prob_ref).cwiseAbs().maxCoeff(), 1e-10);
+}
+
+TEST_F(hmm_test, hidden_state_neg_inf_log_omega) {
+  // A log density of -inf rules out that state for that observation.
+  using stan::math::hmm_hidden_state_prob;
+  Eigen::MatrixXd log_omegas = log_omegas_;
+  log_omegas(0, 4) = -std::numeric_limits<double>::infinity();
+
+  Eigen::MatrixXd prob = hmm_hidden_state_prob(log_omegas, Gamma_, rho_);
+  Eigen::MatrixXd prob_ref
+      = hmm_hidden_state_prob_log_space(log_omegas, Gamma_, rho_);
+
+  EXPECT_EQ(prob(0, 4), 0);
+  EXPECT_TRUE(prob.allFinite());
+  EXPECT_LT((prob - prob_ref).cwiseAbs().maxCoeff(), 1e-12);
+}
+
+TEST(hmm_test_nonstandard, hidden_state_no_states_throws) {
+  // The arguments must be checked before log_omegas is shifted by the
+  // maximum of each column, which is undefined for an empty column.
+  using stan::math::hmm_hidden_state_prob;
+  Eigen::MatrixXd Gamma(2, 2);
+  Gamma << 0.5, 0.5, 0.5, 0.5;
+  Eigen::VectorXd rho(0);
+  Eigen::MatrixXd log_omegas(0, 3);
+  EXPECT_THROW(hmm_hidden_state_prob(log_omegas, Gamma, rho),
+               std::invalid_argument);
+}
+
+TEST_F(hmm_test, hidden_state_no_observations_throws) {
+  using stan::math::hmm_hidden_state_prob;
+  Eigen::MatrixXd log_omegas(n_states_, 0);
+  EXPECT_THROW_MSG(hmm_hidden_state_prob(log_omegas, Gamma_, rho_),
+                   std::invalid_argument,
+                   "hmm_hidden_state_prob: log_omegas has size 0, but must "
+                   "have a non-zero size");
+}