A Primer on Pólya-gamma Random Variables - Part I: Basic Relationships

A recurring nuisance in Bayesian statistics is that the likelihoods we want and the priors we can compute with refuse to cooperate. Bayesian logistic regression is the canonical example: the Bernoulli likelihood with a logit link is not conjugate to a Gaussian prior, the posterior has no closed form, and one ordinarily reaches for approximations (the Laplace approximation, variational inference, expectation propagation) or for sampling (MCMC). We have collectively made peace with this state of affairs, to the point where it is easy to forget there is anything to make peace with.

The Pólya-gamma augmentation of Polson, Scott, and Windle¹ offers a different bargain: introduce auxiliary random variables, chosen so that the awkward likelihood becomes conditionally Gaussian, and the non-conjugate model becomes conjugate again, conditioned on the right things. No approximation is involved; the augmented model is the original model, exactly.

This series builds the trick up from its parts. In this first post we lay out the small collection of identities that make everything work, relating the logistic sigmoid to the hyperbolic cosine, and the hyperbolic cosine to the Laplace transform of the Pólya-gamma distribution. There are no models and no data here; just the moving parts, each verified numerically before we trust it with anything. Part II assembles these parts into Bayesian logistic regression with an exact Gibbs sampler, and Part III swaps the sampler for variational inference and finds a classical method hiding inside.

The Sigmoid, via the Hyperbolic Cosine

Our protagonist is the ubiquitous logistic sigmoid function,

and our first supporting character is the hyperbolic cosine,

Identity 1

The two are related by

This might be immediately obvious for some, but it wasn’t for me. Here is the approach I took to derive it. First, note that

Multiplying by $1 = \frac{e^{\frac{u}{2}}}{e^{\frac{u}{2}}}$ gives

Taking logarithms yields a form that Parts II and III will lean on repeatedly:

The point of the maneuver is that the denominator $2\cosh(\frac{u}{2})$, unlike $1 + e^{-u}$, is something we are about to express as an expectation.

The Pólya-Gamma Distribution

Enter the second supporting character.

The distribution is supported on the positive reals; $b$ acts as a shape parameter and $c$ as a tilting parameter (more on the latter in Part II). The infinite sum means the density has no closed form, which sounds disqualifying for a distribution we intend to build models out of. The saving grace, and the reason the distribution was constructed this way in the first place, is that its Laplace transform has a very pleasant closed form indeed.

Sampling is a solved problem: Polson and Windle devised efficient exact samplers², available in Python via the polyagamma package,

import numpy as np
from polyagamma import random_polyagamma

rng = np.random.default_rng(8888)
omega = random_polyagamma(1, 0, size=32, random_state=rng)

The Laplace Transform

Identity 2

There is a hyperbolic cosine, conveniently in a denominator. To make it match the one in Identity 1, set $b=1$ and reparameterize

which gives

Combining with Identity 1, the sigmoid becomes an expectation with respect to a Pólya-gamma random variable:

This is the identity the entire series stands on. Everything in Parts II and III is a consequence of the fact that the exponent on the right-hand side is quadratic in $u$.

A Mixture of Gaussian Kernels

To see the quadratic structure clearly, give the integrand a name,

so that

It is tempting to read this as a joint density $p(u, \omega)$ with conditional $k$ and marginal $\sigma$. Resist the temptation: $\sigma$ is not integrable in $u$ (it tends to $1$ as $u \to \infty$), and $k(\cdot, \omega)$ is not a normalized density either. What we have is an integral representation of the function $\sigma$ rather than a probability model. The genuinely probabilistic version arrives in Part II, where the observed random variable is the discrete label $y$ and everything normalizes as it should.

Completing the square in the exponent shows that each kernel is a Gaussian bump in $u$, rescaled:

and more precisely³

A draw of $\omega$ thus selects a Gaussian kernel centered at $\frac{1}{2\omega}$ with variance $\frac{1}{\omega}$: small $\omega$ gives a wide bump far to the right, large $\omega$ a narrow bump near the origin. The sigmoid is what you get by averaging these bumps over the $\mathrm{PG}(1, 0)$ distribution.

In code, the kernel is a one-liner,

def kernel(u, omega):
    return 0.5 * np.exp(-0.5 * u * (u * omega - 1.0))

or, in its rescaled-Gaussian form,

from scipy.stats import norm


def kernel(u, omega):

    var = np.reciprocal(omega)

    loc = 0.5 * var
    scale = np.sqrt(var)

    rv = norm(loc=loc, scale=scale)

    return np.sqrt(0.5 * np.pi) * scale * np.exp(0.5**2 * loc) * rv.pdf(u)

Monte Carlo Verification

If the representation is correct, averaging kernels at sampled values of $\omega$ should reproduce the sigmoid:

which is two more lines,

omega = random_polyagamma(1, 0, size=1024, random_state=rng)
sigma_hat = kernel(u[:, None], omega).mean(axis=-1)

The average hugs the sigmoid, as promised. It is a modest kind of magic: a function famous for ruining conjugacy turns out to be an average of the most conjugacy-friendly objects in existence.

What’s Next

Nothing above involved data, parameters, or a posterior. In Part II we attach this representation to an actual model: the observed variable becomes the binary label $y$, the kernel’s quadratic exponent delivers conditional Gaussian-Gaussian conjugacy, and the resulting conditional posteriors give an exact Gibbs sampler for Bayesian logistic regression. In Part III we run variational inference in the same augmented model and watch the classical Jaakkola-Jordan bound emerge from the ELBO, together with an explanation of why that bound is tight exactly where it is.

Links and Further Readings

• Papers:
  • The original augmentation paper (Polson et al., 2013)¹
  • Efficient sampling of Pólya-gamma variates (Windle et al., 2014)²
• Blog posts:
  • Pólya-Gamma Augmentation by G. Gundersen
• Code:
  • polyagamma: the actively maintained Python package used above
  • pypolyagamma: S. Linderman’s Cython port of BayesLogit, J. Windle’s original R implementation

Cite as:

@article{tiao2021polyagammabasic,
  title   = "{A} {P}rimer on {P}ólya-gamma {R}andom {V}ariables - {P}art I: {B}asic {R}elationships",
  author  = "Tiao, Louis C",
  journal = "tiao.io",
  year    = "2021",
  url     = "https://tiao.io/post/polya-gamma-basic-relationships/"
}

To receive updates on more posts like this, follow me on Twitter and GitHub!