Truncated Student t is extremely slow

mathlad · June 5, 2023, 7:10pm

Hello devs,

NUTS sampler is extremely slow with Truncated Student t distribution.

In my case, using Truncated Student t for observation is 4 times slower than Student t. Is it due to tensorflow probability which I had to install specifically for Truncated Student t distribution? Is this a known issue? Thanks!

martinjankowiak · June 5, 2023, 7:48pm

computing log_probs of truncated distributions often involves nasty special functions and so is expected to be slow in many cases. 4x doesn’t sound that bad really

mathlad · June 5, 2023, 7:57pm

Why do we need tensorflow probability for Truncated Student t, but not for Truncated Normal?

martinjankowiak · June 5, 2023, 8:47pm

the warning basically tells you that the incomplete beta function is needed to compute inverse cdfs:

github.com

pyro-ppl/numpyro/blob/4e3c50eff074451ddf909478371aa4dacb18f498/numpyro/distributions/util.py#L435


      
                  from jax.scipy.special import betainc as betainc_fn
          
          
    dtype = jnp.result_type(float)
              return betainc_fn(
                  jnp.array(a, dtype=dtype),
                  jnp.array(b, dtype=dtype),
                  jnp.array(x, dtype=dtype),
              )
          
          

          
def betaincinv(a, b, y):
              try:
                  from tensorflow_probability.substrates.jax.math import special as tfp_special
              except ImportError as e:
                  raise ImportError(
                      "Please install `tensorflow_probability>=0.18` for betaincinv."
                  ) from e
          
          
    dtype = jnp.result_type(float)
              return tfp_special.betaincinv(
                  jnp.array(a, dtype=dtype),