Safe, Scalable, and Accurate Bayes Posterior Sampling for Large-Data Generalized Linear Mixed Models

1. Problem Statement

The paper addresses the computational challenge of performing Bayesian inference for Generalized Linear Mixed Models (GLMMs) on large datasets.

• Context: GLMMs extend Generalized Linear Models (GLMs) to handle grouped data with varying regression coefficients (random effects). Bayesian estimation is preferred for its uncertainty quantification but is computationally expensive due to intractable marginal likelihoods requiring numerical integration over latent variables.
• Limitations of Existing Methods:
  • Standard MCMC: Scales poorly as the number of groups ($n$) increases because it requires sampling both parameters and latent variables, leading to $O(n)$ or worse complexity.
  • Stochastic Gradient MCMC (SGMCMC): Methods like Stochastic Gradient Langevin Dynamics (SGLD) offer scalability via data subsampling but fail in the GLMM context.
    • Instability: GLMMs involve parameters constrained to specific domains (e.g., variance/covariance matrices must be positive definite). Standard SGLD often relies on smooth re-parameterizations (e.g., $\theta = \log \sigma$) to handle these constraints. The authors demonstrate that this approach leads to divergent Markov chains because the gradient of the log-posterior grows super-linearly, causing numerical instability and "blow-up" of the chain.
    • Bias: Even when stable, standard SGLD suffers from irreducible bias in posterior variance estimation due to subsampling noise, often underestimating the true variance of random effects.

2. Methodology

The authors propose a novel algorithm called Stochastic Mirror Langevin Dynamics (SMLD) tailored for constrained sampling in GLMMs, coupled with a rigorous post-processing step.

A. Stochastic Mirror Langevin Dynamics (SMLD)

Instead of re-parameterizing constrained variables into unconstrained space (which causes instability), SMLD operates directly on the constrained space using Mirror Descent principles.

• Mirror Map: A barrier function $\phi(\theta)$ (e.g., $-\log \det(\Omega)$ for covariance matrices) is chosen such that it approaches infinity at the boundary of the feasible set.
• Dual Space Dynamics: The algorithm simulates a stochastic differential equation (SDE) in the dual space defined by the gradient of the mirror map ($\nabla \phi$).
• Update Rule: The discrete update involves a stochastic gradient estimator computed on a mini-batch of data, mapped through the Hessian of the mirror function. This ensures the iterates remain within the feasible set (ergodicity) without requiring Metropolis-Hastings corrections.
• Gradient Estimation: Since the marginal log-likelihood of GLMMs is intractable, the gradient is estimated using Monte Carlo samples of the random effects (latent variables) derived from Fisher's identity.

B. Theoretical Guarantees

The paper provides an explicit Wasserstein distance ($W_2$) error bound between the algorithmic approximation and the true posterior.

• The error bound depends on the step size $\epsilon$, the number of groups $n$, and the dimension $d$.
• The analysis shows that with a step size $\epsilon = O(n^{-(1+\delta)})$, the distance decays as $O(n^{-\delta})$.

C. Post-Processing for Variance Correction

A critical contribution is a non-intrusive post-processing step to correct the bias in posterior variance estimation caused by subsampling.

• Problem: The noise introduced by stochastic gradients inflates the variance of the Markov chain samples compared to the true posterior.
• Solution: The authors derive a Lyapunov equation relating the invariant distribution of the SMLD chain to the true posterior curvature (Hessian of the log-likelihood at the MLE).
  • They estimate the "total covariance" of the noise term using the data and Monte Carlo samples.
  • They solve the Lyapunov equation to recover the true posterior curvature ($H^{-1}$).
  • Correction: The raw SMLD samples are rescaled using the estimated curvature to yield an asymptotically order-wise correct estimate of the posterior variance. This eliminates the need for complex control variates.

3. Key Contributions

1. Algorithmic Safety: Demonstrated that standard SGLD with smooth re-parameterization is unsafe for GLMMs due to transient Markov chains. Proposed SMLD as a stable, ergodic alternative for constrained parameters (specifically covariance matrices).
2. Novel Algorithm: Introduced the first stochastic subsampling analogue of Mirror Langevin Dynamics (SMLD) for Bayesian inference, generalizing SGLD to constrained domains.
3. Variance Correction: Developed a rigorous, non-intrusive post-processing method that corrects the irreducible bias in posterior variance estimation caused by data subsampling, providing accurate uncertainty quantification.
4. Theoretical Analysis: Provided explicit Wasserstein error bounds for SMLD in the context of large-data GLMMs, quantifying the trade-off between step size, sample size, and approximation error.

4. Experimental Results

The method was evaluated through simulated experiments and a real-world case study.

• Toy Model: In a simple normal variance estimation task, standard SGLD diverged (blow-up), while SMLD remained stable. The post-processed SMLD samples matched the true posterior mean and covariance exactly, whereas uncorrected samples showed large bias.
• Simulated GLMMs:
  • Tested on binary response GLMMs with $n=10^4$ and $n=10^5$ groups.
  • Convergence: SMLD mean estimates converged quickly to true values.
  • Variance: Uncorrected SMLD overestimated variance; the proposed correction aligned the variance estimates with those of a full-batch Gibbs sampler (where convergence was achievable).
  • Efficiency: SMLD achieved accurate estimates with significantly less computational time than full-batch MCMC.
• Real-World Data (Breast Cancer Survivors):
  • Applied to a longitudinal study of pain trajectories ($n=2,528$ patients, 33k observations).
  • The corrected SMLD produced credible intervals for fixed effects that were significantly narrower and more precise than uncorrected versions.
  • The correction revealed that while cohort-level pain decreased significantly, individual patient variability (random effects) was substantial. The correction was crucial for calibrating Bayesian p-values.

5. Significance

This work bridges a critical gap in scalable Bayesian inference for hierarchical models.

• Reliability: It solves the "safety" problem of applying stochastic gradient methods to constrained parameters, which is ubiquitous in mixed models (variance components).
• Accuracy: By correcting the variance bias, it ensures that uncertainty quantification (credible intervals) is trustworthy, which is often the primary goal of Bayesian analysis.
• Scalability: It enables the fitting of complex GLMMs on datasets with hundreds of thousands of groups, a regime previously inaccessible to standard MCMC and unstable for SGLD.
• Practicality: The post-processing step is modular and does not require changing the core sampling loop, making it easy to integrate into existing inference pipelines.

In summary, the paper provides a safe, scalable, and statistically accurate framework for Bayesian GLMMs, moving beyond the limitations of current stochastic gradient methods by leveraging mirror dynamics and asymptotic variance correction.