Neural Generalized Mixed-Effects Models

TL;DR

This paper introduces Neural Generalized Mixed-Effects Models (NGMM), enhancing traditional GLMMs with neural networks to better capture complex, nonlinear relationships in hierarchical data.

Contribution

The paper proposes NGMM, a flexible neural network-based extension of GLMMs, along with an efficient, differentiable optimization method for fitting the model to data.

Findings

NGMM outperforms GLMMs on nonlinear covariate-response data.

NGMM achieves better results than prior methods on real-world datasets.

The approximation error of the likelihood decays at a Gaussian-tail rate.

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.