Neural Generalized Mixed-Effects Models

Abstract: Generalized linear mixed-effects models (GLMMs) are widely used to analyze grouped and hierarchical data. In a GLMM, each response is assumed to follow an exponential-family distribution where the natural parameter is given by a linear function of observed covariates and a latent group-specific random effect. Since exact marginalization over the random effects is typically intractable, model parameters are estimated by maximizing an approximate marginal likelihood. In this paper, we replace the linear function with neural networks. The result is a more flexible model, the neural generalized mixed-effects model (NGMM), which captures complex relationships between covariates and responses. To fit NGMM to data, we introduce an efficient optimization procedure that maximizes the approximate marginal likelihood and is differentiable with respect to network parameters. We show that the approximation error of our objective decays at a Gaussian-tail rate in a user-chosen parameter. On synthetic data, NGMM improves over GLMMs when covariate-response relationships are nonlinear, and on real-world datasets it outperforms prior methods. Finally, we analyze a large dataset of student proficiency to demonstrate how NGMM can be extended to more complex latent-variable models.

• The paper introduces NGMM, a novel framework that parameterizes both fixed and random effects using neural networks for enhanced modeling flexibility.
• The methodology leverages an ODE-based marginal likelihood approximation to enable reliable inference in non-conjugate settings across arbitrary exponential family likelihoods.
• Empirical evaluations on synthetic and real-world datasets show NGMM's superior predictive performance, especially for novel groups, compared to traditional models.

Neural Generalized Mixed-Effects Models: An Authoritative Summary

Introduction and Motivation

Mixed-effects modeling is a core statistical paradigm for analyzing data hierarchically structured into groups, inherently capturing between-group and within-group variation via a decomposition into fixed and random effects. The generalized linear mixed-effects model (GLMM) extends linear mixed models to non-Gaussian responses and is standard across numerous scientific domains. However, a linear mapping from covariates to responses, as imposed by GLMMs, is suboptimal for high-dimensional or nonlinear data structures, which are ubiquitous in modern applications.

This paper introduces the Neural Generalized Mixed-Effects Model (NGMM), in which the fixed and random effects are parameterized by neural networks, unifying the flexibility of deep models with the statistical rigor and interpretability of classical mixed-effects frameworks. Unlike prior work, NGMM does not rely on special cases with conjugacy (Gaussian or Gamma-Poisson settings) but supports arbitrary exponential family likelihoods and random effects through a differentiable objective derived from a numerically stable reformulation of the marginal likelihood. This paradigm enables direct optimization of the marginal (rather than conditional) likelihood in non-conjugate settings, which is critical for reliable inference and prediction in novel groups.

Model Specification

NGMM replaces the linear predictors in traditional GLMMs with neural networks. For a grouped dataset, the response $y_{ij}$ for individual $i$ in group $j$ is modeled as:

$$y_{ij} \sim \text{ExpFam} \left( \eta^{-1}(f_\theta(x_{ij}) + g_\psi(z_{j}) \gamma_j) \right),\quad \gamma_j \sim \mathcal{N}(0,I_k)$$

where $f_\theta$ and $g_\psi$ are neural networks parameterized by $\theta$ and $\psi$.

This generalizes random intercept and random slope models. The neural architectures enable NGMM to capture highly nonlinear and non-additive relationships between covariates and response variables.

(Figure 1)

Figure 1: Illustration of the NGMM model architecture, depicting the flow from covariates through neural networks to the group-dependent predictive distribution.

Marginal Likelihood Approximation and Optimization

The central computational challenge is the intractability of the marginal likelihood, resulting from integration over the random effects. NGMM introduces a novel approach by exploiting the orthogonal invariance of the Gaussian prior. The marginalization integral is reduced from a high-dimensional to a one-dimensional problem using a change of variables that aligns with the direction of the random effects’ influence.

The resulting objective is expressed as a non-stiff ODE in the integration variable, yielding a differentiable approximation to the marginal log-likelihood. This formulation is compatible with automatic differentiation, facilitating end-to-end optimization of all neural parameters via stochastic gradient descent. The authors analyze the truncation error of this approximation, proving that the discrepancy between the truncated and exact objectives decays at a Gaussian tail rate as the integration domain expands.

Figure 2: Truncation error and its bound, showing that the approximation error decreases at a Gaussian tail rate as truncation width increases.

The algorithm scales to minibatches over groups and supports further subsampling within groups, enabling deployment on large, structured datasets.

Theoretical Guarantees

The authors provide a rigorous analysis of statistical and computational errors:

• Statistical Error: The empirical risk minimizer based on the ODE-truncated objective is shown to be consistent and asymptotically normal under standard conditions, via M-estimation theory for clustered data.
• Truncation Error: The tail probability for omitting mass outside the integration limits decays as $\exp(-M^2/2)$, and the difference in likelihoods is correspondingly controlled. In the limiting case, the estimator converges to the exact maximum likelihood solution as the truncation parameter grows.

Empirical Evaluation

Synthetic Data

On synthetic grouped data with known nonlinear ground truth, NGMM consistently outperforms GLMMs when the covariate–response relationship is nonlinear (measured by negative log-likelihood, RMSE, and deviance). Under linear data-generating processes, NGMM reduces to a linear model and matches or slightly underperforms GLMMs.

Figure 3: Simulation results; NGMM achieves consistently lower error than GLMM in nonlinear settings, with all values below zero favoring NGMM.

Real-World Datasets

Airbnb: Gaussian & Poisson Outcomes

On the Airbnb listings dataset, NGMM matches specialized neural LMMs for Gaussian outcomes and surpasses Poisson-Gamma NN models for count data, particularly for prediction in new (previously unseen) host groups.

Figure 4: NGMM achieves comparable results to specialized models for Gaussian responses and outperforms all competitors for Poisson responses on Airbnb.

RxRx1 Microscopy: Binary Outcomes

For RxRx1, NGMM demonstrates its capacity to handle binary responses and high-cardinality group structure (microscope field). It delivers lower log-loss and higher accuracy, particularly for novel groups, compared to both closed-form and variational inference baselines. Embedding and one-hot baselines show overfitting to training groups.

Figure 5: In RxRx1, NGMM better approximates marginal log-likelihood and achieves superior prediction for new groups compared to all baselines.

PISA: Latent Variable Mixed Models

The flexibility and extensibility of NGMM are highlighted in the PISA 2012 study, modeling student proficiency as a latent variable influenced by background and school-level effects. NGMM outperforms linear mixed effect and fixed effect neural baselines for all three domains, with the greatest gains observed for prediction in previously observed schools.

Figure 6: NGMM outperforms all baselines in predicting math proficiency across six countries; random effects yield strong gains for familiar schools.

Methodological Comparisons

• Versus Conjugate/Closed-form Neural Mixed Models: NGMM works seamlessly for non-conjugate settings (e.g., logistic and probit), unlike approaches relying on identity/link function conjugacy.
• Versus BBVI: Empirically, the tailored ODE-based marginal likelihood gives superior accuracy and log-loss compared to black-box variational inference in both observed and new group regimes.

Properties and Implications

• Generalization to New Groups: For Gaussian or symmetric binary links, the Bayes-optimal NGMM predictor for new groups coincides with the fixed effects mean—a generalization property inherited from linear mixed models.
• Overfitting and Group Encoding: NGMM mitigates overfitting that plagues methods using direct group encoding (one-hot or embeddings), especially evident in tasks with shifting group structure.
• Latent Variable Integration: The flexibility to incorporate latent variables and heteroskedastic random effects points to NGMM's utility for complex hierarchical data, as seen in the extension to the Rasch IRT model for PISA.

Limitations and Future Directions

The current formulation assumes isotropic Gaussian random effects for both computational tractability and analytic reasons. Extension to more general, potentially non-Gaussian or structured priors would further expand applicability. More broadly, the integration of modern variational and MCMC inference with the NGMM ODE approximation suggests future research for scalability and robustness.

Conclusion

NGMM represents a statistically sound, computationally efficient framework that unifies the expressivity of neural networks with classical mixed-effects modeling. It addresses the limitations of prior work by enabling direct marginal likelihood optimization for arbitrary exponential family likelihoods and supporting complex, hierarchically grouped, nonlinear data. The combination of theoretical guarantees and consistent empirical superiority demonstrates NGMM’s suitability for high-stakes applications in science and engineering where correct uncertainty quantification and generalization to novel groups are essential. The paradigm also sets a template for future research at the intersection of deep learning and hierarchical probabilistic modeling.