Scalable Krylov Subspace Methods for Generalized Mixed Effects Models with Crossed Random Effects

TL;DR

The paper introduces scalable Krylov subspace techniques for generalized mixed effects models with high-dimensional crossed random effects, addressing slow log-determinant evaluations and linear solves that plague Cholesky-based methods. It develops preconditioned CG and SLQ strategies, analyzes convergence through extremal eigenvalues and effective condition numbers, and introduces multiple predictive-variance estimators including stochastic diagonal and Lanczos-based approaches. Empirical results on simulated and real-world data show dramatic speedups (up to ~10,000x) and improved numerical stability, with predictive distributions matching those from exact Cholesky methods. The methods are implemented in GPBoost, providing practical, open-source tools for scalable GLMM inference in large-scale applications.

Abstract

Mixed-effects models are widely used to model data with hierarchical grouping structures and high-cardinality categorical predictor variables. However, for high-dimensional crossed random effects, sparse Cholesky decompositions, the current standard approach, can become prohibitively slow. In this work, we present Krylov subspace-based methods that address these computational bottlenecks and analyze them both theoretically and empirically. In particular, we derive new results on the convergence and accuracy of the preconditioned stochastic Lanczos quadrature and conjugate gradient methods for mixed-effects models, and we develop scalable methods for calculating predictive variances. In experiments with simulated and real-world data, the proposed methods yield speedups by factors of up to about 10,000 and are numerically more stable than Cholesky-based computations as implemented in state-of-the-art packages such as lme4 and glmmTMB. Our methodology is available in the open-source C++ software library GPBoost, with accompanying high-level Python and R packages.

Figures (14)

• Figure 1: Accuracy (in standard deviations; top row), wall-clock time (in seconds; bottom row, left axis), and number of CG iterations (bottom row, right axis) for different preconditioners and varying numbers of random vectors $t$ in the SLQ method for calculating log-marginal likelihoods for a Gaussian likelihood and balanced and unbalanced random effects designs.
• Figure 2: Accuracy-runtime comparison of different methods for predictive variances. The number of random vectors $s$ and the Lanczos rank $k$ are annotated in the plot.
• Figure 3: Average wall clock times (s) for parameter estimation and different $m$. Simulated data follows either a Gaussian or a Bernoulli likelihood.
• Figure 4: Estimated variance parameters. The red rhombi represent means. The dashed lines indicate the true values. For the Bernoulli likelihood, the estimates for lme4 are not computed due to excessively long runtimes.
• Figure 5: RMSE for predictive means and log score (LS) for probabilistic predictions for Gaussian and Bernoulli likelihoods.

Theorems & Definitions (27)

• Proposition 3.1
• Proposition 3.2
• Proposition 3.3
• Theorem 4.1
• Lemma 4.1
• Lemma 4.2
• Theorem 4.2
• Theorem 4.3
• Theorem 4.4
• Corollary 4.1