The Best Gamma-Poisson Parameterization imo.

I’ve been studying Bayesian models for count data using Gamma-Poisson distributions. The objective was simple, using auxiliary features, use a neural network to predict the parameters of the Gamma distribution, and then use the Gamma distribution as a prior for the count data. The issue arises when trying to actually implement it. Any naive reparameterization I tried (using torch.log, adding epsilons everywhere, normalizing the count data, etc.) didn’t work, and training a neural network always resulted in exploding gradients. I needed to understand the Gamma-Poisson distribution better.

First of all, documentation on Gamma distributions are messy, not only are there different reparameterizations (scale vs. rate), sometimes people use the same symbols for different parameterizations and don’t tell you which ones they’re using! Not to mention some references just get it wrong and the formulas are not even correct (Look at this. Here $\alpha$ and $\beta$ should be switched. yikes).

Let’s develop a good parameterization from first principles. Turns out, besides using scale or rate, there’s a quirky third reparameterization that’s far superior to all others (in my opinion). Not sure if it’s been done before, but I haven’t seen it anywhere, so I’m posting it here. To get at it, let’s first write down the pdf of the Gamma-Poisson distribution. It’s just a mixture of Poisson distributions with parameter $\lambda$ drawn from a Gamma distribution with parameters (shape = $\alpha$) and (rate = $\beta$):

Notice here that the integral looks almost like a Gamma distribution, but with different (shape = $x + \alpha$) and (rate = $1 + \beta$) parameters. It’s also missing a few constants. Since the integral of a probability distribution must be 1, we can just multiply and divide by the missing constants $\frac{\Gamma(x + \alpha)}{(1 + \beta)^{x + \alpha}}$ to get:

Rearranging a little bit, we get:

Notice that I separated the term with $\beta$ and $1 + \beta$ into two terms. I did this to show that the Gamma-Poisson is very closely related to the so-called Negative Binomial distribution. The Negative Binomial distribution is defined as:

We can see that the Negative Binomial distribution and Gamma-Poisson distribution are related as $\alpha = r$ and $\frac{\beta}{1 + \beta} = p$. This is a way better reparameterization than the usual scale or rate parameterization. As we will see, it allows us to use Binary Cross Entropy to compute the log-likelihood, a function which has been heavily optimized in many libraries and is very stable. So, lets parameterize it in terms of logits $z$, and $\sigma(z)$ is the sigmoid of $z$:

Then we can rewrite the Gamma-Poisson distribution as:

The log-likelihood of this distribution is:

Notice that the two last terms looks like Binary Cross Entropy (BCE)!

Specifically, after further reparameterizing the shape $\alpha$ as log shape $\log(\alpha)$, the input logits, targets, and weights to the BCE are:

• Logits: $z$

• Targets: $\frac{\alpha}{x + \alpha} = \text{Softmax}(\log(\alpha), \log(x))[0]$

• Weights: $x + \alpha$

We can just use the Binary Cross Entropy loss function BCEWithLogitsLoss for half of the Gamma-Poisson likelihood computation:

Meanwhile, the other half can just be computed using torch.lgamma() functions. The full Gamma-Poisson likelihood computation is:

def gamma_poisson_loss(self, log_shape, logits, value):
    eps = 1e-8
    concentration = torch.exp(log_shape)
    
    log_value = torch.log(value + eps)
    targets = torch.softmax(
        torch.stack(
            [log_shape, log_value], dim=-1
        ), 
        dim=-1
    )[..., 0]
    
    weights = value + concentration
    
    log_likelihood = (
        torch.lgamma(concentration + value + eps)
        - torch.lgamma(concentration + eps)
        - weights * F.binary_cross_entropy_with_logits(
            input = logits,
            target = targets,
            reduction = 'none'
        )
    )
    
    return -log_likelihood.mean()

It’s also nice because to get the rate, we can just exponentiate the logits:

In conclusion, we parameterized the Gamma-Poisson distribution using logits and log shape. These parameters are unrestricted, so we can ask a neural network to output logits $z$ and $\log(\alpha)$ for the Gamma-Poisson distribution. We then use Binary Cross Entropy for parts the likelihood computation, and torch.lgamma() functions for the rest. Voila! No more exploding gradients!

-Randy