Stan Coding Corner: O(N) Change-Point Program with Clever Forward-Backward Calculation

It’s so much fun to work in open source. Luke Wiklendt sent along this improved code for a change-point model calculation in Stan. With N data points in the time series, the version in the manual is O(N2), whereas the improved version is O(N). In practice, Luke says

[the new code] results in a dramatic speed up. On my machine 10k samples of the original took 8m53s and the new one took 33s.

Now that’s the kind of improvement I can get behind, both in theory and in practice. As soon as I understood his code, I slapped my head and yelled “D’oh!”. Of course there’s a linear algorithm, because it can be encoded as a simple state-space HMM. Everyone (who spent 25 years working on natural language processing and speech recognition) knows the HMM forward-backward algorithm is linear in N.

Quadratic version from Stan manual v2.9

Here’s my original version.

transformed parameters {
  vector[T] lp;
  lp <- rep_vector(log_unif, T);
  for (s in 1:T)
    for (t in 1:T)
      lp[s] <- lp[s] + poisson_log(D[t], if_else(t < s, e, l));
}

Linear version from Luke Wiklendt

And here's Luke's improved version, which he prefaces by saying

Given a vector v and two indices i and j, the sum from i to j is performed by v[i] + v[i+1] + ... + v[j-1] + v[j]. It can also be performed by precalculating the cumulative sum of v (call it cv), and then the sum from i to j is cv[j] - cv[i].

That's what the engineers call "dynamic programming."

transformed parameters {
    vector[T+1] lp_e;
    vector[T+1] lp_l;
    vector[T] lp;

    lp_e[1] <- 0;
    lp_l[1] <- 0;
    for (t in 1:T) {
        lp_e[t + 1] <- lp_e[t] + poisson_log(D[t], e);
        lp_l[t + 1] <- lp_l[t] + poisson_log(D[t], l);
    }

    lp <- rep_vector(log_unif + lp_l[T + 1], T) + head(lp_e, T) - head(lp_l, T);
}

The key for me was realizing this is just the forward-backward algorithm of HMMs. The forward values are lp_e and the backward values are lp_l[T+1] - lp_l, so lp is the complete forward+backward values (plus log_unif for the prior).

Thanks, Luke!

3 thoughts on “Stan Coding Corner: O(N) Change-Point Program with Clever Forward-Backward Calculation

Leave a Reply

Your email address will not be published. Required fields are marked *