From bc338df4470fd554dffce6a8c7874bb36a345d98 Mon Sep 17 00:00:00 2001 From: Ryan Curtin Date: Tue, 16 Mar 2021 15:54:18 -0400 Subject: [PATCH] Use LogSumExp() computations throughout (where possible). --- src/mlpack/methods/hmm/hmm_impl.hpp | 96 ++++++++++++++--------------- 1 file changed, 48 insertions(+), 48 deletions(-) diff --git a/src/mlpack/methods/hmm/hmm_impl.hpp b/src/mlpack/methods/hmm/hmm_impl.hpp index 0449484250..05b4de2e3f 100644 --- a/src/mlpack/methods/hmm/hmm_impl.hpp +++ b/src/mlpack/methods/hmm/hmm_impl.hpp @@ -149,8 +149,8 @@ double HMM::Train(const std::vector& dataSeq) backwardLog, logScales); // Add to estimate of initial probability for state j. - for (size_t j = 0; j < logTransition.n_cols; ++j) - newLogInitial[j] = math::LogAdd(newLogInitial[j], stateLogProb(j, 0)); + math::LogSumExp(stateLogProb.unsafe_col(0), + newLogInitial); // Define a variable to store the value of log-probability for data. arma::mat logProbs(dataSeq[seq].n_cols, logTransition.n_rows); @@ -163,7 +163,6 @@ double HMM::Train(const std::vector& dataSeq) emission[i].LogProbability(dataSeq[seq], alias); } - // Now re-estimate the parameters. This is the M-step. // pi_i = sum_d ((1 / P(seq[d])) sum_t (f(i, 0) b(i, 0)) // T_ij = sum_d ((1 / P(seq[d])) sum_t (f(i, t) T_ij E_i(seq[d][t]) b(i, @@ -172,24 +171,31 @@ double HMM::Train(const std::vector& dataSeq) // We store the new estimates in a different matrix. for (size_t t = 0; t < dataSeq[seq].n_cols; ++t) { - for (size_t j = 0; j < logTransition.n_cols; ++j) + // Assemble temporary vector that's used in log-sum computation. + if (t < dataSeq[seq].n_cols - 1) { - if (t < dataSeq[seq].n_cols - 1) - { - // Estimate of T_ij (probability of transition from state j to state - // i). We postpone multiplication of the old T_ij until later. - for (size_t i = 0; i < logTransition.n_rows; i++) - { - newLogTransition(i, j) = math::LogAdd(newLogTransition(i, j), - forwardLog(j, t) + backwardLog(i, t + 1) + logProbs(t + 1, i) - - logScales[t + 1]); - } - } + // This term is the same across all states, so compute it once and + // cache it. + const arma::vec tmp = backwardLog.col(t + 1) + + logProbs.row(t + 1).t() - logScales[t + 1]; + arma::vec output; + math::LogSumExp(tmp, output); - // Add to list of emission observations, for Distribution::Train(). - emissionList.col(sumTime) = dataSeq[seq].col(t); - emissionProb[j][sumTime] = exp(stateLogProb(j, t)); + for (size_t j = 0; j < logTransition.n_cols; ++j) + { + // Compute the estimate of T_ij (probability of transition from + // state j to state i). We postpone multiplication of the old T_ij + // until later. + arma::vec tmp2 = output + forwardLog(j, t); + arma::vec alias = newLogTransition.unsafe_col(j); + math::LogSumExp(tmp2, alias); + } } + + // Add to list of emission observations, for Distribution::Train(). + for (size_t j = 0; j < logTransition.n_cols; ++j) + emissionProb[j][sumTime] = exp(stateLogProb(j, t)); + emissionList.col(sumTime) = dataSeq[seq].col(t); sumTime++; } } @@ -712,14 +718,12 @@ arma::vec HMM::ForwardAtT0(const arma::vec& emissionLogProb, // P(X_k | o_{1:k}) for all possible states X_k, for each time point k. ConvertToLogSpace(); - arma::vec forwardLogProb(logTransition.n_rows); - forwardLogProb.fill(-std::numeric_limits::infinity()); // The first entry in the forward algorithm uses the initial state // probabilities. Note that MATLAB assumes that the starting state (at // t = -1) is state 0; this is not our assumption here. To force that // behavior, you could append a single starting state to every single data // sequence and that should produce results in line with MATLAB. - forwardLogProb = logInitial + emissionLogProb; + arma::vec forwardLogProb = logInitial + emissionLogProb; // Normalize probability. logScales = math::AccuLog(forwardLogProb); @@ -741,17 +745,16 @@ arma::vec HMM::ForwardAtTn(const arma::vec& emissionLogProb, // Our goal is to calculate the forward probabilities: // P(X_k | o_{1:k}) for all possible states X_k, for each time point k. - arma::vec forwardLogProb(logTransition.n_rows); - forwardLogProb.fill(-std::numeric_limits::infinity()); - // Now compute the probabilities for each successive observation. - for (size_t state = 0; state < logTransition.n_rows; state++) - { - // The forward probability of state j at time t is the sum over all - // states of the probability of the previous state transitioning to - // the current state and emitting the given observation. - arma::vec tmp = prevForwardLogProb + logTransition.row(state).t(); - forwardLogProb(state) = math::AccuLog(tmp) + emissionLogProb(state); - } + // The forward probability of state j at time t is the sum over all states of + // the probability of the previous state transitioning to the current state + // and emitting the given observation. To do this computation in log-space, + // we can use LogSumExp(). + arma::vec forwardLogProb; + arma::mat tmp = logTransition + repmat(prevForwardLogProb.t(), + logTransition.n_rows, 1); + math::LogSumExp(tmp, forwardLogProb); + forwardLogProb += emissionLogProb; + // Normalize probability. logScales = math::AccuLog(forwardLogProb); if (std::isfinite(logScales)) @@ -809,23 +812,20 @@ void HMM::Backward(const arma::mat& dataSeq, // Now step backwards through all other observations. for (size_t t = dataSeq.n_cols - 2; t + 1 > 0; t--) { - for (size_t j = 0; j < logTransition.n_rows; j++) - { - // The backward probability of state j at time t is the sum over all state - // of the probability of the next state having been a transition from the - // current state multiplied by the probability of each of those states - // emitting the given observation. - for (size_t state = 0; state < logTransition.n_rows; state++) - { - backwardLogProb(j, t) = math::LogAdd(backwardLogProb(j, t), - logTransition(state, j) + backwardLogProb(state, t + 1) - + logProbs(t + 1, state)); - } + // The backward probability of state j at time t is the sum over all + // states of the probability of the next state having been a transition + // from the current state multiplied by the probability of each of those + // states emitting the given observation. To compute this in log-space, we + // can use LogSumExpT(). + const arma::mat tmp = logTransition + + repmat(backwardLogProb.col(t + 1), 1, logTransition.n_cols) + + repmat(logProbs.row(t + 1).t(), 1, logTransition.n_cols); + arma::vec alias = backwardLogProb.unsafe_col(t); + math::LogSumExpT(tmp, alias); - // Normalize by the weights from the forward algorithm. - if (std::isfinite(logScales[t + 1])) - backwardLogProb(j, t) -= logScales[t + 1]; - } + // Normalize by the weights from the forward algorithm. + if (std::isfinite(logScales[t + 1])) + backwardLogProb.col(t) -= logScales[t + 1]; } }