Numpyro: HMM forward algorithm - 20x difference in performance between two comparable approaches

[svilupp]

This issue links to the following thread:
https://forum.pyro.ai/t/intuition-for-the-difference-between-the-two-hmm-tutorials-forward-algo-vs-marginalization/2775

The following two tutorials for HMMs cover two alternative solutions via the forward algorithm (one explicit, one implicit):

1. HMM Example semi-supervised
2. Enumating HMM

Context:

• Using simulate_data function from [1] with the following parameters:

args = parser.parse_args("-n 2000 --num-words 10 --num-categories 3 --num-supervised 0 --num-unsupervised 2000".split(’ '))

In other words, 1 series of 2000 observed "words" with transmission matrix 3x3 and emission matrix 3x10 (unsupervised = categories are unknown)

• Estimate model_1 on unsupervised_words from [1] without any changes

• Adjust model_1 from [2] to solve the same task (output categorical dist + align shapes), vis. model_1_alt below

  def model_1_alt(unsupervised_words):

    num_categories=args.num_categories
    num_words=args.num_words
  
    emission_prior=jnp.repeat(0.1, num_words)
    transition_prior = jnp.ones(num_categories)
  
    probs_x = numpyro.sample(
        "probs_x", dist.Dirichlet(
        jnp.broadcast_to(transition_prior, (num_categories, num_categories))).to_event(1)
    )
    probs_y = numpyro.sample(
        "probs_y",dist.Dirichlet(
        jnp.broadcast_to(emission_prior, (num_categories, num_words))).to_event(1)
    )   
  
  def transition_fn(carry, y):
        x_prev = carry
        x = numpyro.sample("x", dist.Categorical(probs_x[x_prev]))  
        y = numpyro.sample("y", dist.Categorical(probs_y[x]), obs=y)
        return x, None
  
   x_init = 0
   scan(transition_fn, (x_init, 0), unsupervised_words)
  

Expected behaviour: Both approaches would achieve similar results within comparable time (eg, +-100%) due to slightly different computations

Actual behaviour: Approach with model_1 from [1] runs c. 20x faster than model_1_alt based on [2] (130s vs 2,700s, respectively; time measures the whole fitting procedure)

Setup:

• Seeds: Simulation done with random.PRNGKey(1); mcmc is done with random.PRNGKey(2)
• ran on CPU on AWS EC2 m5.xlarge
• Python (3.7.8), Numpyro installed from master (0.6.0), Jax (0.2.12)

[fehiepsi]

I did a bit of benchmarking and it seems to me that this is expected for parallel-scan algorithm

import jax
import jax.numpy as jnp
from jax.scipy.special import logsumexp
import numpyro; numpyro.set_platform("gpu")  # "cpu"
import funsor; funsor.set_backend("jax")


def logmatmulexp(x, y):
    x_shift = x.max(-1, keepdims=True)
    y_shift = y.max(-2, keepdims=True)
    return jnp.log(jnp.exp(x - x_shift) @ jnp.exp(y - y_shift)) + x_shift + y_shift, None


@jax.jit
def sequential(x_init, xs):
    o, _ = jax.lax.scan(logmatmulexp, xs[0], xs[1:])
    o = logmatmulexp(jnp.expand_dims(x_init, -2), o)[0]
    return logsumexp(o.squeeze(-2), -1)


@jax.jit
def forward(x_init, xs):
    o, _ = jax.lax.scan(logmatmulexp, jnp.expand_dims(x_init, -2), xs)
    return logsumexp(o.squeeze(-2), -1)


@jax.jit
def parallel(x_init, xs):
    batch_shape = xs.shape[:-3]
    state_dim = xs.shape[-1]
    while xs.shape[-3] > 1:
        time = xs.shape[-3]
        even_time = time // 2 * 2
        even_part = xs[..., :even_time, :, :]
        a_b = even_part.reshape(batch_shape + (even_time // 2, 2, state_dim, state_dim))
        a, b = a_b[..., 0, :, :], a_b[..., 1, :, :]
        contracted = logmatmulexp(a, b)[0]
        if time > even_time:
            contracted = jnp.concatenate((contracted, xs[..., -1:, :, :]), axis=-3)
        xs = contracted
    o = logmatmulexp(jnp.expand_dims(x_init, -2), xs.squeeze(-3))[0]
    return logsumexp(o.squeeze(-2), -1)


@jax.jit
def funsor_scan(x_init, xs):
    trans = funsor.Tensor(xs)["time", "prev", "curr"]
    o = funsor.sum_product.sequential_sum_product(
        funsor.ops.logaddexp,
        funsor.ops.add,
        trans,
        funsor.Variable("time", funsor.Bint[xs.shape[0]]),
        {"prev": "curr"})
    o = logmatmulexp(jnp.expand_dims(x_init, -2), o.data)[0]
    return logsumexp(o.squeeze(-2), -1)


dim = 3
x = jax.random.normal(jax.random.PRNGKey(0), (2000, dim, dim))
x_init = jax.random.normal(jax.random.PRNGKey(1), (dim,))
sequential(x_init, x)
parallel(x_init, x)
forward(x_init, x)
funsor_scan(x_init, x)
%timeit y = sequential(x_init, x).copy()
%timeit y = parallel(x_init, x).copy()
%timeit y = forward(x_init, x).copy()
%timeit y = funsor_scan(x_init, x).copy()

which returns

# on GPU
sequential
52 ms ± 96 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
parallel
328 µs ± 1.81 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
forward
44.3 ms ± 68.7 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
funsor_scan
442 µs ± 2.74 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# on CPU
sequential
1.42 ms ± 32.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
parallel
1.41 ms ± 31.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
forward
222 µs ± 4.76 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
funsor_scan
1.43 ms ± 35.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

This shows the huge benefit of the parallel-scan algorithm (100x faster) on GPU. On CPU, the parallel-scan implementation is equivalent to sequential implementation but 7x slower than the forward algorithm, which exploits the fact that v @ (A @ B @ C @ D) can be evaluated faster by computing v @ A @ B @ C @ D from left to right (computing a sequence of matrix-matrix products is slower than a sequence of vector-matrix products). This is an interesting performance issue but I'm not sure how to improve the performance if we stick with the parallel-scan algorithm. Probably we can use the forward algorithm on the CPU but setting up the pattern to use forward algorithm in more complicated models would be pretty tricky. Do you really need this speed up or just want to understand the reason for the difference?

cc @eb8680 @fritzo who helped me identify the issue, thanks!

[svilupp]

Thank you for the detailed answer, fehiepsi!
I've just realized that there was a question for me - yes, I just wanted to understand the reason for the big performance difference, because it seemed odd.

Your simple code was awesome! Until I played with it across different dimensions, I haven't appreciated how big the GPU toll is for some operations!