Use LogSumExp() computations throughout (where possible).

src/mlpack/methods/hmm/hmm_impl.hpp

@@ -149,8 +149,8 @@ double HMM<Distribution>::Train(const std::vector<arma::mat>& 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<arma::vec, true>(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<Distribution>::Train(const std::vector<arma::mat>& 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<Distribution>::Train(const std::vector<arma::mat>& 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<arma::vec, true>(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<Distribution>::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<double>::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<Distribution>::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<double>::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<Distribution>::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<arma::mat, true>(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];
   }
 }