Enabling Mixed Effects Neural Networks for Diverse, Clustered Data Using Monte Carlo Methods

TL;DR

The paper addresses the challenge of modeling clustered data in deep learning by introducing MC-GMENN, a generalized mixed effects neural network trained with Monte Carlo Expectation Maximization. The model combines a fixed-effects neural network f_Omega(x) with cluster-specific random effects B^(l) per clustering feature and class, assuming Gaussian priors and optimizing the marginal likelihood. Training alternates between a Monte Carlo E-step, which samples random effects from their posterior using No-U-Turn Sampling, and an M-step that updates the fixed-effects parameters and variance components, with an additional term to strengthen fixed-effect learning. Empirically, MC-GMENN outperforms existing MENNs in generalization, provides accurate inter-cluster variance quantification, remains time-efficient, and scales to multi-class problems with multiple clustering features, while enabling clear interpretability of cluster effects.

Abstract

Neural networks often assume independence among input data samples, disregarding correlations arising from inherent clustering patterns in real-world datasets (e.g., due to different sites or repeated measurements). Recently, mixed effects neural networks (MENNs) which separate cluster-specific 'random effects' from cluster-invariant 'fixed effects' have been proposed to improve generalization and interpretability for clustered data. However, existing methods only allow for approximate quantification of cluster effects and are limited to regression and binary targets with only one clustering feature. We present MC-GMENN, a novel approach employing Monte Carlo methods to train Generalized Mixed Effects Neural Networks. We empirically demonstrate that MC-GMENN outperforms existing mixed effects deep learning models in terms of generalization performance, time complexity, and quantification of inter-cluster variance. Additionally, MC-GMENN is applicable to a wide range of datasets, including multi-class classification tasks with multiple high-cardinality categorical features. For these datasets, we show that MC-GMENN outperforms conventional encoding and embedding methods, simultaneously offering a principled methodology for interpreting the effects of clustering patterns.

Figures (3)

• Figure 1: Illustration of a generalized mixed effects neural network with MCEM parameter updates for binary classification.
• Figure 2: Comparison of the learned random effects distribution of different mixed effects neural network approaches on a simulated dataset with one clustering feature ($Q=1000$ and $\sigma^2=1.0$).
• Figure 3: Interpretation example for learned random effects distributions on the open_payments dataset. The task is to classify whether payments from manufacturers to physicians are disallowed. The payments are clustered by manufacturer (Q=1460), the device associated with the payment (Q=4365), the drug associated with the payment (Q=2255), and the physician's specialty (Q=513). The dashed lines show estimated random effects for one example payment.