Handle zero-mean Negative Binomial distributions

[dylanhmorris]

Description

The negative binomial distribution is well-defined with mean $\mu = 0$. Equivalently, it is well-defined with $p = 0$ in the probs/total_count parametrization used by NegativeBinomialProbs. In practice, NegativeBinomial2 fails validation with mean 0. More concerningly, NegativeBinomialProbs passes validation with probs = 0 but then behaves incorrectly, yielding nan log prob values. Forward sampling and evaluation of the distribution mean work correctly.

This derives from the fact that NegativeBinomialProbs(probs=jnp.array(0)) and NegativeBinomial2(mean=jnp.array(0)) both end up constructing a GammaPoisson with an inf rate:

numpyro/numpyro/distributions/conjugate.py

Line 483 in a1bcf5e

rate = 1.0 / probs - 1.0

numpyro/numpyro/distributions/conjugate.py

Line 554 in a1bcf5e

rate = concentration / mean

Action

I think this can and should be patched. I am happy to contribute a PR, but would appreciate some discussion on approach. In particular, do we want inf-rate GammaPoisson instances to be acceptable and well-defined? My instinct is yes but I can imagine there might be sharp bits.

Details and reprex

For NegativeBinomial2, the mean is given a constraints.positive, so construction fails with validate_args=True. Note: construction fails before validation with a divide-by-zero error if the mean is an int or float rather than a jax or numpy array.

import jax.numpy as jnp
import numpyro.distributions as dist

dist.NegativeBinomial2(mean = jnp.array(0), concentration = 1, validate_args=True) # fails validation
# ValueError: NegativeBinomial2 distribution got invalid mean parameter.

dist.NegativeBinomial2(mean = 0, concentration = 1, validate_args=True) # divide-by-zero error
# ZeroDivisionError: division by zero

For NegativeBinomialProbs, probs is given constraints.unit_interval. Construction and evaluation of the log prob therefore succeed with probs=jnp.array(0) regardless of whethervalidate_args=True (Note: as in the NegativeBinomial2 case, construction fails with a divide by zero error if the probs is a python int or float). But we get nan log prob values:

dist.NegativeBinomialProbs(total_count=1, probs=0, validate_args=True) # divide by 0 error
# ZeroDivisionError: float division by zero

my_neg_bin = dist.NegativeBinomialProbs(total_count=1, probs=jnp.array(0), validate_args=True) # construction succeeds

my_neg_bin.log_prob(0) # nan log prob; should be 0
# Array(nan, dtype=float32, weak_type=True)

my_neg_bin.log_prob(1) # nan log prob; should be -inf
# Array(nan, dtype=float32, weak_type=True)

Forward sampling and evaluation of the mean behaves correctly. In that case, the underlying GammaPoisson handles the inf rate gracefully:

my_neg_bin.mean # correct mean 
# Array(0., dtype=float32, weak_type=True)

# correct forward sampling
with numpyro.handlers.seed(rng_seed=5):
    samples = numpyro.sample('neg bin samples', my_neg_bin.expand((100,)))

all(samples == 0)
# True

Behavior is similar if we construct a zero-mean NegativeBinomial2 with validate_args=False. This makes sense; either way, we end up with a GammaPoisson with an inf rate.

my_neg_bin2.log_prob(0) # nan log prob, should be 0
# Array(nan, dtype=float32, weak_type=True)

my_neg_bin2.log_prob(1) # nan log prob, should be -inf
# Array(nan, dtype=float32, weak_type=True)

my_neg_bin2.mean # correct mean
# Array(0., dtype=float32, weak_type=True)

with numpyro.handlers.seed(rng_seed=5):
    samples = numpyro.sample('neg bin samples', my_neg_bin2.expand((100,)))

all(samples == 0)
# True

[Qazalbash]

Thanks for the report @dylanhmorris.

For NegativeBinomial2, the mean is given a constraints.positive, so construction fails with validate_args=True. Note: construction fails before validation with a divide-by-zero error if the mean is an int or float rather than a jax or numpy array.

Arguments validation is checked at the end of the __init__ method of each class, when super().__init__(...) is called. The division is happening before this call, so when you have zero integer/float mean, Python raises a division by zero error before it can reach the argument validation step.

numpyro/numpyro/distributions/distribution.py

Lines 277 to 278 in a1bcf5e

	if self._validate_args:
	self.validate_args(strict=False)

numpyro/numpyro/distributions/conjugate.py

Lines 554 to 555 in a1bcf5e

	rate = concentration / mean
	super().__init__(concentration, rate, validate_args=validate_args)

Same case for NegativeBinomialProbs,

numpyro/numpyro/distributions/conjugate.py

Lines 483 to 484 in a1bcf5e

	rate = 1.0 / probs - 1.0
	super().__init__(concentration, rate, validate_args=validate_args)

For NegativeBinomialProbs, probs is given constraints.unit_interval. Construction and evaluation of the log prob therefore succeed with probs=jnp.array(0) regardless of whether validate_args=True (Note: as in the NegativeBinomial2 case, construction fails with a divide by zero error if the probs is a python int or float). But we get nan log prob values:

This is expected as per the implementation, maybe we can use xlogy and xlop1py to make the following lines more stable,

numpyro/numpyro/distributions/conjugate.py

Lines 396 to 397 in a1bcf5e

	+ self.concentration * jnp.log(self.rate)
	- post_value * jnp.log1p(self.rate)

Forward sampling and evaluation of the mean behaves correctly. In that case, the underlying GammaPoisson handles the inf rate gracefully:

Can you show how you constructed the GammaPoisson object?

In particular, do we want inf-rate GammaPoisson instances to be acceptable and well-defined? My instinct is yes but I can imagine there might be sharp bits.

IMO, if inf-rate GammaPoisson has an analytical form independent of any inf term, then it can be implemented as a separate distribution or maybe in the same distribution. One can also use jax.numpy.finfo(jax.numpy.float32).tiny as an equivalent to zero.

[dylanhmorris]

This is expected as per the implementation, maybe we can use xlogy and xlop1py to make the following lines more stable

I don't think that will fix this problem. The nan arises from the fact + self.concentration * jnp.log(self.rate) evaluates to inf, - post_value * jnp.log1p(self.rate) evaluates to -inf, and jnp.inf - jnp.inf is undefined. But + xlogy(self.concentration, self.rate) and - xlog1py(post_value, self.rate) still evaluate to inf and -inf, respectively, because concentration > 0 and therefore post_value = self.concentration + value > 0.

[dylanhmorris]

The docs give the parametrization used here for the PMF (with a typo; have opened #2194).

If $X \sim \mathrm{GammaPoisson}(\alpha, \lambda)$, then:

$$p_{X}(k) = \frac{\lambda^\alpha}{(\alpha + k)(1+\lambda)^{\alpha + k}\mathrm{B}(\alpha, k + 1)} $$

Rewrite this as:

$$p_{X}(k) = \left(\frac{\lambda}{1+\lambda}\right)^\alpha \left[\frac{1}{(\alpha + k)(1+\lambda)^{k}\mathrm{B}(\alpha, k + 1)}\right] $$

The second term in that product should be computable with inf rates $\lambda$, but the first is problematic. Idea: compute that term as:

$$\left(1 + \frac{1}{\lambda}\right)^{-\alpha}$$

In log space:

def log_rate_ratio_reg(rate, concentration):
    return xlogy(concentration, rate) - xlogy(concentration, 1 + rate)

def log_rate_ratio_new(rate, concentration):
    return xlogy(-concentration, 1 + 1 / rate)

The two give similar values for finite rate, but the second returns 0 for inf rate while the first returns nan:

log_rate_ratio_reg(9.7, 3.5)
# Array(-0.34341288, dtype=float32, weak_type=True)

log_rate_ratio_new(9.7, 3.5)
# Array(-0.34341252, dtype=float32, weak_type=True)

log_rate_ratio_reg(jnp.inf, 3.5)
# Array(nan, dtype=float32, weak_type=True)

log_rate_ratio_new(jnp.inf, 3.5)
# Array(-0., dtype=float32, weak_type=True)

So we could consider writing a new log prob for GammaPoisson that looks something like this:

def log_prob(self, value):
    return(
        xlogy(-self.concentration, 1 + 1 / self.rate) -
        jnp.log(self.concentration + value) -
        xlogy(value, 1 + self.rate) -
        betaln(self.concentration, value + 1)
    )

In some very non-systematic testing, this is similar to the existing implementation for non-inf rates, with the desired behavior at inf rates. But I am aware that there might be sharp bits here.

[Qazalbash]

These examples are pretty convincing; we can implement them. Do you want to open a PR?

[dylanhmorris]

These examples are pretty convincing; we can implement them. Do you want to open a PR?

Sure, can do!