hmm_hidden_state returns NaN

[charlesm93]

Description

Several users have reported that hmm_hidden_state returns nan, in cases where the probability of the hidden state being in one particular state goes to 1. There seems to be some numerical stability issue. Computing the log unnormalized probability first and then using softmax solves the issue in the reported problems.

This issue is raised on the Stan forum, notably here. A proposed coding of the method in Stan is given here.

Example

See discussion on Stan forum.

Expected Output

API doesn't change but function returns 1 or 0, as it does when manually writing the code in Stan with a log_sum_exp.

Current Version:

v4.3.0

[passiflorai]

Hi, here is the data and associated code to reproduce the issue. Let me know if any problems with code/data.
stan_question.tar.gz

[charlesm93]

@passiflorai I ran your code and can confirm hmm_hidden_state indeed returns an error.

I tried building a simpler example which triggers the error, see math/test/unit/math/prim/prob/hmm_hidden_state_prob_test.cpp on branch bugfix/issue-2677-hmm_hidden_state. The last test creates a situation where the probability of being in one state is 1.

The numerical issue only arises when the log density of the observational model is large in magnitude. I tried setting log_omega to 1,000 or -1,000 and got NaN. Set all the elements of log_omega to 1 and there is no problem. The same observation can be made for other unit tests, where the prob of being in a state isn't 0 or 1. Hmmm... I'm not sure the issue is related to having the hidden probability go to 1.

EDIT: Let's take a closer look at your example and extract the expectation value of log_omega. In R

log_omega <- matrix(colMeans(fit$draws('log_omega')), ncol = 2,
                    byrow = FALSE)

At first glance, the values look reasonable, even towards the end where the probability goes to 1. But along the way there are a few large values in magnitude.

> log_omega[order(log_omega)][1:5]
[1] -53.82944 -42.92558 -35.78769 -34.17773 -33.51622

In the C++ code (hmm_hidden_state_prob.hpp), all these get exponentiated

Eigen::MatrixXd omegas = value_of(log_omegas).array().exp();

So we get very small elements (e.g. exp(-53) = 9.60268e-24). Would this be enough to cause a numerical issue when calculating the alphas?

for (int n = 0; n < n_transitions; ++n)
    alphas.col(n + 1)
        = omegas.col(n + 1).cwiseProduct(Gamma_dbl.transpose() * alphas.col(n));