Polya-Gamma Sampling For Bayesian Logistic Regression

Why Augment a Logistic Regression?

In Bayesian logistic regression, the log likelihood is not quadratic in the regression coefficients. With a normal prior, that means the coefficient posterior is generally not a normal distribution that can be sampled by a direct conjugate update. Polya-Gamma augmentation introduces an auxiliary variable for each binomial observation. Conditional on those variables, the likelihood becomes quadratic in the coefficients, giving a multivariate normal full conditional.

The method is useful when a Gibbs sampler with a direct coefficient update is desirable. It does not remove the need to assess the resulting chain: the auxiliary variables and coefficients can still be strongly dependent, and computation may be demanding for large or poorly conditioned models.

Model and Assumptions

Suppose there are $n$ independent observations. For row $i$, $y_i$ is the number of events among $m_i$ trials, where $m_i$ is a positive integer and $0 \leq y_i \leq m_i$. Let $\mathbf{x}_i$ be a fixed $p$-vector of predictors, and let $\boldsymbol{\beta}$ be the coefficient vector. The binomial logistic-regression model is

y_i \mid \boldsymbol{\beta} \sim \operatorname{Binomial}(m_i,p_i), \qquad \operatorname{logit}(p_i)=\psi_i=\mathbf{x}_i^\mathsf{T}\boldsymbol{\beta}, \qquad i=1,\ldots,n.

The observations are assumed conditionally independent given $\boldsymbol{\beta}$, and the trials within each row share probability $p_i$. The design vectors are treated as known. Use a proper multivariate normal prior

\boldsymbol{\beta}\sim N(\mathbf{b}_0,\mathbf{B}_0),

where $\mathbf{B}_0$ is positive definite. Define the $n \times p$ design matrix $\mathbf{X}$, whose $i$th row is $\mathbf{x}_i^\mathsf{T}$, and let $\boldsymbol{\psi}=\mathbf{X}\boldsymbol{\beta}$. Rows with zero trials, if present in a data file, contain no binomial information and should be removed before applying the updates below; the formulas here use $m_i>0$.

Definition: For shape $b>0$ and tilt $c\in\mathbb{R}$, a Polya-Gamma random variable $\omega\sim PG(b,c)$ is a positive random variable obtained by exponentially tilting the $PG(b,0)$ distribution by $\exp(-\omega c^2/2)$, with the normalizing constant chosen to make a probability distribution. The base distribution is characterized by $$ E\{\exp(-t\omega)\}=\cosh^{-b}\!\left(\sqrt{t/2}\right),\qquad t\geq 0. $$

The Augmentation Identity

Write the binomial likelihood for one row, omitting the binomial coefficient because it does not depend on $\boldsymbol{\beta}$:

p_i^{y_i}(1-p_i)^{m_i-y_i} = \frac{\exp(y_i\psi_i)}{(1+\exp(\psi_i))^{m_i}}.

Set $\kappa_i=y_i-m_i/2$. Using $1+\exp(\psi)=2\exp(\psi/2)\cosh(\psi/2)$, the same likelihood becomes

\frac{\exp(y_i\psi_i)}{(1+\exp(\psi_i))^{m_i}} = 2^{-m_i}\exp(\kappa_i\psi_i)\cosh^{-m_i}(\psi_i/2).

By the defining Laplace transform of $PG(m_i,0)$, the last factor has the mixture representation

\cosh^{-m_i}(\psi_i/2) = E_{\omega_i\sim PG(m_i,0)} \left\{\exp(-\omega_i\psi_i^2/2)\right\}.

Thus an auxiliary variable $\omega_i$ can be introduced so that the augmented likelihood, up to factors not involving $\boldsymbol{\beta}$ or $\omega_i$, is proportional to

\exp\left(\kappa_i\psi_i-\frac{\omega_i\psi_i^2}{2}\right) \,p_{PG(m_i,0)}(\omega_i).

This representation gives both full conditional distributions used by the sampler. The coefficient-dependent part is quadratic in $\boldsymbol{\beta}$, and the auxiliary-variable conditional is the exponentially tilted Polya-Gamma distribution.

Worked Example: Rewriting One Binomial Likelihood

Consider a hypothetical row with $m_i=4$ trials and $y_i=3$ events. Its logit is $\psi_i$, and $\kappa_i=y_i-m_i/2=3-2=1$. Ignoring the binomial coefficient, the likelihood is

\frac{\exp(3\psi_i)}{(1+\exp(\psi_i))^4} = 2^{-4}\exp(\psi_i)\cosh^{-4}(\psi_i/2).

The augmentation identity replaces the hyperbolic-cosine factor by an expectation over $\omega_i\sim PG(4,0)$:

\frac{\exp(3\psi_i)}{(1+\exp(\psi_i))^4} = \frac{1}{16} E\left\{\exp\left(\psi_i-\frac{\omega_i\psi_i^2}{2}\right)\right\}.

The term involving the logit inside this expectation is a linear term $\psi_i$ minus a quadratic term $\omega_i\psi_i^2/2$. This is the structure that leads to a normal coefficient update.

Full Conditional Distributions

Conditional on the current coefficient vector, the auxiliary variables are independent, with

\omega_i\mid\boldsymbol{\beta},\mathbf{y} \sim PG(m_i,\psi_i), \qquad \psi_i=\mathbf{x}_i^\mathsf{T}\boldsymbol{\beta}.

For the coefficient update, define $\boldsymbol{\kappa}$ as the vector of $\kappa_i$ values and $\boldsymbol{\Omega}=\operatorname{diag}(\omega_1,\ldots,\omega_n)$. Because every $m_i>0$, each $\omega_i>0$ with probability one. Combining the augmented likelihood with the normal prior gives the coefficient-dependent kernel

\exp\left[ \boldsymbol{\kappa}^\mathsf{T}\mathbf{X}\boldsymbol{\beta} -\frac{1}{2}\boldsymbol{\beta}^\mathsf{T} \mathbf{X}^\mathsf{T}\boldsymbol{\Omega}\mathbf{X}\boldsymbol{\beta} -\frac{1}{2}(\boldsymbol{\beta}-\mathbf{b}_0)^\mathsf{T} \mathbf{B}_0^{-1}(\boldsymbol{\beta}-\mathbf{b}_0) \right].

Collecting the quadratic and linear terms yields a multivariate normal full conditional:

Result: Given $\boldsymbol{\omega}$ and the data, $$ \boldsymbol{\beta}\mid\boldsymbol{\omega},\mathbf{y} \sim N(\mathbf{V}\mathbf{h},\mathbf{V}), \qquad \mathbf{V}=(\mathbf{X}^\mathsf{T}\boldsymbol{\Omega}\mathbf{X}+\mathbf{B}_0^{-1})^{-1}, \qquad \mathbf{h}=\mathbf{X}^\mathsf{T}\boldsymbol{\kappa}+\mathbf{B}_0^{-1}\mathbf{b}_0. $$

Why the Coefficient Conditional Is Normal

The prior precision $\mathbf{B}_0^{-1}$ is positive definite. Since $\boldsymbol{\Omega}$ has positive diagonal entries, $\mathbf{X}^\mathsf{T}\boldsymbol{\Omega}\mathbf{X}$ is positive semidefinite. Their sum is therefore positive definite and invertible, even if the columns of $\mathbf{X}$ are linearly dependent.

To verify the normal form, expand the log kernel and discard terms constant in $\boldsymbol{\beta}$. It becomes

-\frac{1}{2}\boldsymbol{\beta}^\mathsf{T}\mathbf{V}^{-1}\boldsymbol{\beta} +\boldsymbol{\beta}^\mathsf{T}\mathbf{h}.

Completing the square gives

-\frac{1}{2} (\boldsymbol{\beta}-\mathbf{V}\mathbf{h})^\mathsf{T} \mathbf{V}^{-1}(\boldsymbol{\beta}-\mathbf{V}\mathbf{h}) +\frac{1}{2}\mathbf{h}^\mathsf{T}\mathbf{V}\mathbf{h}.

The final term does not depend on $\boldsymbol{\beta}$. The remaining kernel is that of a normal distribution with mean $\mathbf{V}\mathbf{h}$ and covariance $\mathbf{V}$, proving the stated update.

Worked Example: A Scalar Intercept Update

Suppose a hypothetical intercept-only model has two rows: $(m_1,y_1)=(4,4)$ and $(m_2,y_2)=(2,0)$. Use the prior $\beta\sim N(0,4)$. At a particular Gibbs iteration, suppose the auxiliary-variable draws are $\omega_1=0.8$ and $\omega_2=0.5$. These are illustrative draws, not observed data. Here $\kappa_1=4-2=2$, $\kappa_2=0-1=-1$, and both design values equal 1.

The prior precision is $1/4=0.25$. The conditional precision and linear term are

V^{-1}=0.8+0.5+0.25=1.55, \qquad h=2+(-1)+0.25(0)=1.

Therefore the conditional variance and mean are

V=\frac{1}{1.55}=0.6452, \qquad Vh=\frac{1}{1.55}=0.6452.

The coefficient draw is from $N(0.6452,0.6452)$, where the second argument is the variance. The calculation is conditional on these particular auxiliary draws; a new pair of auxiliary draws at the next iteration generally produces a different conditional distribution.

Worked Example: Updating Intercept and Treatment Effect

Consider a hypothetical two-row trial model with an intercept and a treatment indicator. The rows have $(m_1,y_1)=(2,1)$ for control and $(m_2,y_2)=(2,2)$ for treatment. The design matrix is

\mathbf{X}= \begin{pmatrix} 1&0\\ 1&1 \end{pmatrix}, \qquad \boldsymbol{\beta}= \begin{pmatrix} \beta_0\\ \beta_1 \end{pmatrix}.

Use the prior $N(\mathbf{0},\mathbf{I})$, and suppose the current auxiliary draws are $\omega_1=\omega_2=0.5$. The row values are $\kappa_1=1-1=0$ and $\kappa_2=2-1=1$. The precision is

\mathbf{V}^{-1} = \mathbf{X}^\mathsf{T}\boldsymbol{\Omega}\mathbf{X}+\mathbf{I} = \begin{pmatrix}1&0.5\\0.5&0.5\end{pmatrix} + \begin{pmatrix}1&0\\0&1\end{pmatrix} = \begin{pmatrix}2&0.5\\0.5&1.5\end{pmatrix}.

The linear term is $\mathbf{h}=\mathbf{X}^\mathsf{T}\boldsymbol{\kappa}=(1,1)^\mathsf{T}$. The determinant of the precision matrix is $2(1.5)-(0.5)^2=2.75$, so

\mathbf{V} = \frac{1}{2.75} \begin{pmatrix}1.5&-0.5\\-0.5&2\end{pmatrix}, \qquad \mathbf{V}\mathbf{h} = \begin{pmatrix}0.3636\\0.5455\end{pmatrix}.

Thus the joint conditional update is a bivariate normal draw with the displayed mean and covariance $\mathbf{V}$. The second coefficient is the treatment log-odds contrast in this model; the conditional draw is not by itself a final estimate of that contrast.

Gibbs Sampling Steps

The augmented model gives a two-step Gibbs sampler. At iteration $t$, use the current coefficients to update the auxiliary variables, then use those variables to update all coefficients jointly.

Draw auxiliary variables.
For each row, calculate $\psi_i^{(t)}=\mathbf{x}_i^\mathsf{T}\boldsymbol{\beta}^{(t)}$, then draw independently $\omega_i^{(t+1)}\sim PG(m_i,\psi_i^{(t)})$.

Draw coefficients jointly.
Form $\mathbf{V}$ and $\mathbf{h}$ using $\boldsymbol{\Omega}^{(t+1)}$, then draw $\boldsymbol{\beta}^{(t+1)}\sim N(\mathbf{V}\mathbf{h},\mathbf{V})$.

The joint coefficient draw matters when coefficients are correlated: it updates their dependence as part of one multivariate draw. As with other Gibbs samplers, initial values and retained draws should be accompanied by appropriate chain diagnostics and Monte Carlo error assessment.

Diagnostics, Alternatives, and Limitations

Several checks help detect implementation errors. Verify that every $m_i$ is positive and every $y_i$ lies between 0 and $m_i$; check that the computed $\kappa_i$ equals $y_i-m_i/2$; and confirm that $\boldsymbol{\Omega}$ has positive entries and the coefficient precision matrix is symmetric positive definite. Simulated $\omega_i$ values must be positive. For a small test case, compare the implemented mean and covariance with the formulas above and check that simulated coefficient draws have plausible scales under the chosen prior.

Confusing precision and covariance. The matrix $\mathbf{X}^\mathsf{T}\boldsymbol{\Omega}\mathbf{X}+\mathbf{B}_0^{-1}$ is the conditional precision; its inverse is the conditional covariance.
Using the wrong tilt. The auxiliary draw uses $PG(m_i,\psi_i)$, with the current linear predictor as its tilt, not $PG(m_i,0)$ at every iteration.
Leaving zero-trial rows in the update. The stated distribution is defined here for positive shape. A row with $m_i=0$ contributes no likelihood information and should be removed.
Assuming conjugacy means independent draws. Each full conditional is directly sampleable, but successive Gibbs iterations can still be correlated. Mixing and effective information in the chain require assessment.

Compared with Metropolis-within-Gibbs for logistic trial models, this method avoids choosing a proposal scale and an accept-or-reject step for the coefficients. It also provides a direct normal coefficient update. Hamiltonian Monte Carlo offers a different route to exploring a continuous posterior and can be effective in higher-dimensional models, but it requires its own tuning and diagnostics. Polya-Gamma sampling is not automatically faster: the auxiliary-variable draws add computation, and strong posterior dependence can limit mixing. Its normal update also does not make the marginal posterior for $\boldsymbol{\beta}$ normal; normality holds only conditional on the auxiliary variables.

The method relies on the binomial logistic likelihood and the stated conditional-independence structure. A proper normal prior makes the coefficient precision well-defined even with redundant predictors, but it does not resolve a poorly chosen model or guarantee good practical mixing. In an implementation, the Polya-Gamma random-number generator must accurately sample the required shape and tilt distributions.

Key takeaway: With $m_i>0$, the Polya-Gamma identity makes each auxiliary-variable conditional $PG(m_i,\psi_i)$ and makes the joint coefficient conditional multivariate normal. The Gibbs sampler alternates between these two exact full conditional updates.

Check Your Understanding

Answer each question, then compare your work with the solutions.

For a row with $m_i=6$ and $y_i=2$, calculate $\kappa_i$. What is the conditional distribution of $\omega_i$ given the current linear predictor $\psi_i$?
In an intercept-only model with one row, $m=2$, $y=2$, prior $\beta\sim N(0,1)$, and current auxiliary draw $\omega=0.5$, find the conditional mean and variance of $\beta$.
Why does a proper normal prior ensure that the coefficient conditional precision is invertible, even if the design matrix has redundant columns?
What is wrong with using $\omega_i\sim PG(m_i,0)$ at every iteration, regardless of the current coefficient draw?

Solutions. (1) $\kappa_i=2-6/2=-1$, and $\omega_i\mid\boldsymbol{\beta},\mathbf{y}\sim PG(6,\psi_i)$. (2) The conditional precision is $\omega+1=1.5$, and the linear term is $\kappa+1(0)=2-1=1$. Thus the variance is $1/1.5=0.6667$ and the mean is $1/1.5=0.6667$. (3) The data contribution $\mathbf{X}^\mathsf{T}\boldsymbol{\Omega}\mathbf{X}$ is positive semidefinite, while the prior precision is positive definite; their sum is positive definite and invertible. (4) The correct tilt depends on the current linear predictor: the full conditional is $PG(m_i,\psi_i)$. Repeatedly using the untilted distribution would not implement the augmented model’s Gibbs conditional.