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

1 Introduction

Clustering patterns are evident in data across various domains, such as medicine [Cafri et al., 2019], ecology [Harrison et al., 2018], or e-commerce [Fei et al., 2021]. For instance, in product return forecasting, transaction samples are naturally grouped by customer, product, brand, or geographic location. These clusters can often number in the thousands, with each cluster containing only a small number of samples.

In Deep Neural Networks (DNNs), clustering information is commonly treated as an additional categorical feature, often integrated through numeric encoding (e.g., onehot encoding) or embeddings [Hancock and Khoshgoftaar, 2020; Borisov et al., 2021]. While these approaches improve predictive performance compared to ignoring cluster information, they may encounter issues of overfitting, overparameterization, and scalability when dealing with highcardinality categorical features [Simchoni and Rosset, 2021]. Furthermore, the models blend cluster information with other features, making it challenging to interpret the specific effects of cluster membership accurately.

In the statistics community, generalized linear mixed models (GLMMs) are well-established for handling clustered data [McCulloch, 2003; Pinheiro and Chao, 2006; Agresti, 2012]. Recently, there has been growing interest in integrating GLMMs with deep learning [Xiong et al., 2019; Simchoni and Rosset, 2023; Nguyen et al., 2023]. Mixed Effects Neural Networks (MENNs) are partially Bayesian models that use fixed effects DNNs and incorporate clustering features separately as probabilistic random effects.

Existing MENN approaches have demonstrated improved predictive performance and interpretability over conventional encoding and embedding approaches. The main challenge in training MENNs is that the negative log-likelihood loss for classification has no closed-form expression. While Markov Chain Monte Carlo (MCMC) methods are common for traditional GLMMs [McCulloch, 1997; Archila, 2016], modern Bayesian Neural Networks are more frequently trained using variational inference (VI) due to its time efficiency [Blei et al., 2017; Jospin et al., 2022]. Consequently, all existing MENN approaches rely on approximate methods like VI, although MCMC could provide an exact quantification of the inter-cluster variance [Jospin et al., 2022]. This limits the interpretability and thus invalidates the main reason for using mixed effects instead of conventional approaches.

A previously underappreciated fact is that, unlike fully Bayesian neural networks, MENNs only need to sample the parameters of the random effects, which changes the way scalability considerations need to be made. Moreover, modern MCMC methods, particularly the No-U-Turn Sampler (NUTS) [Hoffman and Gelman, 2011] greatly speed up convergence of MCMC algorithms compared to the time when MCMC for GLMMs was introduced [McCulloch, 1997].

Based on these insights we propose MC-GMENN, an approach to train generalized mixed effects neural networks by combining state-of-the-art MCMC methods and deep learning in an Expectation Maximization (EM) framework (Section 2). In Section 3 we demonstrate that our approach:

• outperforms existing mixed effects deep learning approaches in performance, time efficiency, and intercluster variance quantification (Section 3.1).
• scales well to a variety of datasets with high dimensionalities, no. of clustering features, no. of classes, cardinalities, and inter-cluster variance constellations (Section 3.2).
• outperforms encoding and embedding approaches on 16 classification benchmark datasets with multiple highcardinality categorical features while providing high interpretability (Section 3.3).

A major factor preventing wider adoption of MENNs is that existing approaches do not apply to classification with multiple high-cardinality clustering features and classes, as we will discuss in Section 4. This paper addresses this gap and, to our knowledge, represents the first empirical demonstration of mixed effects deep learning performance for classification tasks with multiple classes and clustering features.

2 Monte Carlo GeneralizedMixed Effects Neural Networks (MC-GMENN)

In this section, we introduce our generalized model formulation and the proposed Monte Carlo Expectation Maximization (MCEM) training procedure. For a general introduction to GLMMs we refer to Chapter 3 in [Agresti, 2012] and [McCulloch, 2003]. Existing MENNs and differences to our contribution will be discussed in Section 4.

2.1 GeneralizedMixed Effects Neural Networks

Let X ∈ RN×D be the fixed effects design matrix and Y ∈ {0, 1}N×C be a matrix indicating class membership, where N is the number of samples, D is the number of (fixed effects) features, and C is the number of classes. In addition, let Z be a set of random effects design matrices Z(l) ∈ {0, 1}N×Ql with information about cluster membership for L categorical features of cardinalities Q1, ..., QL.

We formulate our GMENN model as:

where fΩ is a neural network parameterized by Ω and B(l) ∈ RQl×C is a matrix with random effect vectors per class for clustering feature l. ϕ is an activation function depending on whether the target is continuous, binary or multiclass. For simplified notation, let B be the set of all random effects vectors.

The model is based on the assumptions of traditional GLMMs [McCulloch, 1997]:

1. The samples yi are conditionally independent given the random effects B and drawn from a distribution py|B in the exponential family suitable to describe the target.
2. The random effects b(11) , ..., b(LC) are assumed to be independent and distributed according to parametric distributions p(11) b , ..., p(LC) b . Most commonly, Normal distributions with zero mean are used for each p(lc) b : b(lc) ∼ N (0, Σ(lc) ).

The model is depicted in Figure 1. By assuming a distribution on the effect of categorical features, the random effects regularize the estimates of cluster effects. Thereby, the amount of regularization (variance parameters) is learned from the data [Sigrist, 2023]. In the case of the most popular random effects model, the random intercept model, Σ simplifies to σ2 I. For simplification, let S be the set of all covariance matrices.

To fit the model with parameters Ω, S, and B, we need to maximize the marginal data likelihood:

2.2 Monte Carlo ExpectationMaximization for GMENN

To evaluate the intractable marginal log-likelihood function, we extend the Monte Carlo Expectation Maximization (MCEM) approach for GLMMs [McCulloch, 1997] 1 to the state-of-the-art of deep learning and MCMC methods. For better readability, we use matrix notation throughout the remainder of the paper. In the EM framework, Y represents the observed data, while B remains unobserved, with unknown variance parameters S. We substitute Equation 2 with a Monte Carlo approximation of its expected value.

E-Step. In each epoch t, we create a function to evaluate the expectation of the log-likelihood of Equation 2 under the current parameter estimates, given the observed data Y:

1 An implementation of the traditional MCEM algorithm for linear mixed effects regression is available at https://www.tensorflow. org/probability/examples/Linear Mixed Effects Models

where

To evaluate this expectation, we would ordinarily need to compute Equation 2, which we aim to avoid. Consequently, we employ Monte Carlo integration to estimate the expectation. We generate K sets of samples B(1) , . . . , B(K) from the conditional distributions p(11) b |Y, . . . , p(LC) b |Y at each epoch t. Among these, the initial R samples serve as burnin. In contrast to existing approaches, we utilize the No-UTurn Sampler (NUTS) [Hoffman and Gelman, 2011] for sampling. NUTS is known to demonstrate remarkable efficiency in traversing complex likelihood surfaces, thus significantly reducing the required number of samples [Hoffman and Gelman, 2011; Monnahan and Kristensen, 2018]. An ablation study demonstrating the benefits of NUTS compared to other samplers can be found in the supplementary material. Moreover, NUTS can be automated to operate without the need for hyperparameters, making it particularly well-suited for enhancing both the time efficiency and user-friendliness of our MCEM framework.

M-step. In the M-step, we update Ω and S using the Monte Carlo estimate of E(t) and gradient descent. Because B is available as MCMC samples, the two terms in Equation 3 can be decoupled:

To ensure that fΩ converges and performs well on its own to predict unseen clusters, we add an additional term to Equation 5a, inspired by [Nguyen et al., 2023]. Hence, the fixed effects loss becomes:

where 0 denotes the random effects set to zero, and λ is a hyperparameter that controls how much emphasis should be placed on the fixed effects.

Convergence of MC-GMENN is determined by early stopping on validation data using Equation 5a or a performance metric. After convergence, estimations of the random effects ˆb(lc) are obtained as the mean of all samples over all epochs. Predictions are obtained using the estimated random effects coefficients in the model (Equation 1).

Important MCEMProperties for Deep Learning. Three key properties make the MCEM procedure exceptionally suitable for combining mixed effects and deep learning: First, the sampling (E-step) is detached from the mini-batch loss evaluation (M-step). In the E-step, the computation of fΩ(X) for the first term of Equation 3 is required only once. Hence, the E-step can be evaluated efficiently without necessitating mini-batches. In the M-step, the random effects B are available as MCMC samples. Hence, no expensive integration procedure limits the speed of the mini-batch gradient descent updates. Consequently, the algorithm scales well to large neural networks, high fixed effects feature dimensionality and sample size. Second, the updates in Equations 5a and 5b can be computed independently. Depending on the specific task, varying update policies and early stopping rules can be used for Equation 5b. In the case of the Gaussian random intercept model, which is the most popular choice, the variance parameters S are updated by setting them to the variance of the current epoch's samples. If exact variance estimation is required, the training can be continued by only updating 5b until the variance parameters converge. Third, Equation 5a aligns with the conventional loss function formulation, facilitating the use of common losses that have a corresponding log-likelihood counterpart such as cross-entropy during the M-step. Furthermore, Equation 5a naturally decomposes into mini-batches, which is not the case for all MENNs.

Hyperparameters and Automation. The additional hyperparameters added by MC-GMENN are the (initial) step size ϵ of NUTS, the no. of samples K, the no. of epochs to use as burn-in R and the fixed effects weight λ. A too large ϵ exponentially increases the probability of rejection, while a too small ϵ is very time-consuming. We found that common step size adaption methods like dual averaging [Hoffman and Gelman, 2011] do not perform well. Instead, we propose to start with a large step size of ϵ0 = 0.1 and divide it by two whenever the acceptance rate gets below 0.001. Due to the ability of NUTS to traverse large distances, one informative sample per epoch can suffice for a good estimation. Moreover, a different configuration than λ = 1 is only required if the fixed effects exhibit significantly slower convergence compared to the random effects. In Section 3, we show that no hyperparameter optimization is necessary to achieve competitive results with our method as it consistently demonstrates strong performance across a diverse range of datasets.

3 Experiments

In all experiments, we compare MC-GMENN to the following conventional methods: Ignoring high-cardinality categorical features (Ignore), one-hot-encoding (OHE), target encoding (TE) [Micci-Barreca, 2001], and entity embeddings (Embedding) [Guo and Berkhahn, 2016]. For each method, we use the same base neural network architecture and optimizer with the same hyperparameters (learning rate, decay, dropout, embedding size), as well as training procedure (epochs, patience, batch size) per experiment. To prove that our method is readily applicable to any dataset, we use the default hyperparameters for MC-GMENN over all experiments: ϵ0 = 0.1,K = 2,R = 1, λ = 1. In line with previous work [Simchoni and Rosset, 2023; Nguyen et al., 2023], we use the area under the ROC curve (AUC) as the performance metric. To be able to compare across different datasets, training time is evaluated relative to the Ignore method. For simulated datasets, we additionally evaluate the ability of the MENN models to learn the underlying random effects distribution. For that, we use the absolute error of the estimated variance components and visually compare the learned distributions. Our framework and all the evaluated models are implemented in TensorFlow. 2 More detailed information about the datasets, hyperparameters, and evaluation setup are provided in the supplementary material.

3.1 Comparison of MC-GMENNwith Related Mixed Effects Approaches

Experimental Setting. To achieve a fair evaluation of our method and existing approaches, we replicate the binary classification experiments of LMMNN, one of the most recent MENN approaches [Simchoni and Rosset, 2023]. Using the data generation method described in their paper, we generate datasets with features nonlinearly related to the target and additional categorical clustering features with specified cardinalities. For the latter, the target varies according to a Normal distribution with specified variance depending on the cluster membership. Nine scenarios are simulated for five iterations with varying Q ∈ {100, 1000, 10000} and σ2 ∈ {0.1, 1.0, 10.0}. The utilized neural network architecture consists of four hidden layers with 100, 50, 25, and 12 neurons with ReLU activation and dropout regularization of 25%. Furthermore, our evaluation incorporates target encoding, MC-GMENN, and ARMED [Nguyen et al., 2023], with the latter representing another recent and noteworthy MENN approach. For ARMED, we use the same hyperparameters as Nguyen et al. [2023] in their simulated experiments.

MC-GMENN Outperforms Existing Mixed Effects Approaches. As can be seen in Table 1, MC-GMENN shows the best overall performance. Compared to the reported results in Simchoni and Rosset [2023], our replication shows very similar performances, training times, and variance quantification (see supplementary material). This indicates that our replication was correct and fair towards LMMNN and strengthens the superiority of our proposed approach. Unexpectedly, our findings indicate that ARMED underperforms in high-cardinality scenarios which will be further elaborated in Section 4.

2 Our code is available at https://github.com/atschalz/mcgmenn

MC-GMENN Precisely Quantifies Inter-Cluster Variance. Figure 2 shows that only our MC-GMENN approach is able to fit the posterior distribution correctly. LMMNN uses a Gaussian quadrature approximation which forces the random effects close to the quadrature points, leading to a bad approximation. At the same time, the training time is very high already, such that refining the approximation by using more quadrature points is infeasible, especially in highcardinality settings. It can also be seen that ARMED, which uses a VI approach, is strongly biased towards the prior variational distribution (σ2 = 0.01). This hinders the correct estimation of larger cluster effects. The result is a worse posterior estimation and ultimately a worse performance. In contrast, MC-GMENN provides an unbiased estimate of the integral and thus is able to learn the underlying distribution correctly.

MC-GMENN is Time-Efficient Compared to Existing Mixed Effects Approaches. The training time comparison in Table 1 shows that MC-GMENN is more efficient than the competitors. For ARMED, high-cardinality features lead to a parameter explosion greatly increasing the training time, as will be discussed in Section 4. In LMMNN, the loss function contains matrix inversions that are evaluated in mini-batches making it inefficient. In contrast, the MCEM procedure of MC-GMENN separates the expensive part in the E-step such that the mini-batch updates in the M-step are as time-efficient as with a regular neural network.

3.2 Scalability toMulti-Class Datasets with Multiple Clustering Features

Experimental Setting. To evaluate the performance of MC-GMENN when applied to datasets with multiple clustering features and classes, we extend the data generation algorithm of Simchoni and Rosset [2023] to multi-class classification. As will be discussed in Section 4, LMMNN and ARMED are not applicable to these datasets. To demonstrate the scalability and versatility of MC-GMENN, we simulate datasets for one realistic base scenario with five repetitions and vary this scenario according to relevant data properties. The base scenario is defined as a multi-class classification task with N = 100, 000, D = 10, C = 5, and three clustering features with Q1 = 1000, Q2 = 10, Q3 = 1000. The variance is assumed to be constant per class with σ2 1 = 0.0001, σ2 2 = 0.5, σ2 3 = 0.5. Hence, there is one categorical feature with no impact and two with medium impact on the target, one of which has low and one has high cardinality. Further scenarios are simulated with five repetitions each by varying the base scenario as described in Table 2. The same base neural network architecture as in Subsection 3.1 is used.

MC-GMENN Consistently Matches or Outperforms Encoding and Embedding Approaches. As can be seen in Table 2, MC-GMENN shows the best average performance across all scenarios. Even in scenarios where it is not the single best model, the relative difference to the best model is low. Moreover, Table 3 shows that MC-GMENN quantifies the inter-cluster variance in multi-class settings with low mean absolute differences to the true variances.

MC-GMENN Scales to Various Data Scenarios. The superior performance proves, that MC-GMENN is applicable to different data dimensionality or clustering variance structures. As can be expected, the training time relative to the Ignore condition increases for more complex scenarios. It can be observed that the relative increase in training time is higher for datasets with more complex random effects structures, such as the 100 classes scenario and the high-cardinality scenario. The relative training time for 1000 features is almost the same as for 10 features and does not greatly increase with 1 million samples. This demonstrates that, as discussed in Subsection 2.2, MC-GMENN scales well to large datasets.

3.3 Application to Diverse Real-World Datasets

Experimental Setting. To evaluate the performance of MC-GMENN on real-world classification tasks, we use the same datasets as in Pargent et al. [2022], which is, to our knowledge, the largest benchmark study on the treatment of high-cardinality categorical features. We use the same preprocessing as Pargent et al. [2022], where Q = 10 is used as a threshold to define high-cardinality features. Low-cardinality categorical features are encoded using OHE. Other categorical features are treated as clustering features. More details on the datasets and preprocessing can be found in the supplementary material and in the paper of Pargent et al. [2022]. For all approaches, we use the neural network architecture implemented in the AutoML framework AutoGluon-tabular with its default hyperparameters [Erickson et al., 2020].

MC-GMENN Consistently Matches or Outperforms Encoding and Embedding Approaches. The results in Table 4 demonstrate strong performance of MC-GMENN across various real-world datasets. As in the previous subsection, MC-GMENN particularly excels on datasets with numerous high-cardinality clustering features and classes. Even for datasets where MC-GMENN does not achieve the top rank, the performance gap to the best model remains small. For some datasets, the Ignore method ranks among the topperforming models. The fact that MC-GMENN exhibits comparable performance on most of these datasets highlights its strong capability in discerning irrelevant clustering features, which is an important property for tabular deep learning [Grinsztajn et al., 2022]. In conclusion, MC-GMENN more consistently achieves high performance than competitive approaches.

MC-GMENN Enables Interpretability of Cluster Effects. Figure 3 illustrates the interpretability of MC-GMENN on an example. The estimated random effects distribution can be used to assess global clustering feature importance. The estimated variances indicate that all four categorical features have clusters with impact, while manufacturer and drug have generally the most widespread distributions and thus potentially stronger effects. As the random effects are linear, they provide white box access to the model behavior w.r.t. to the categorical features and can be used for local interpretability. For the illustrated sample, the payment is classified as disallowed. The model has estimated a high random effect on the logits for the manufacturer making the payment. Hence, the classification is greatly influenced by the fact that the payment is made by this particular manufacturer.

4 RelatedWork

In this section, we discuss related work with a focus on mixed effects deep learning approaches for classification. We refer to [Prokhorenkova et al., 2018; Pargent et al., 2022] for the treatment of categorical data in general and to [Hancock and Khoshgoftaar, 2020; Borisov et al., 2021; Huang et al., 2020] for examples of how categorical data is typically treated in deep learning. Existing MENN approaches are different from each other mainly in terms of the training procedure and their applicability, as summarized in Table 5.

Xiong et al. [2019] propose MeNets, a MENN relying on variational EM with stochastic gradient descent that was originally proposed for regression. Notably, our EM framework is fundamentally different, as Xiong et al. [2019] update the fixed effects parameters in the E-step and rely on the expensive inversion of large covariance matrices to update the random effects. In contrast, we use MCMC to obtain random effect samples in the E-step and perform all parameter updates in the M-step. Nguyen et al. [2023] showed that for binary classification, the approach is less efficient and less performant than LMMNN and ARMED, therefore we did not include it in our evaluation.

Tran et al. [2020] propose DeepGLMM, a MENN for panel data trained with Bayesian variational approximation. Although theoretically widely applicable, the approach was only evaluated in a limited setting on one real-world dataset without comparison to other deep learning methods. Furthermore, the authors of recent MENN approaches were not able to use the implementation due its computational and conceptual complexity [Simchoni and Rosset, 2023; Nguyen et al., 2023]. In contrast, we demonstrated that MC-GMENN is easily applicable to diverse datasets even with the default hyperparameters.

Simchoni and Rosset [2021] focus on regression and develop a plug-in loss function to train linear mixed model neural networks (LMMNN). To this end, an important aspect is the decomposability of the loss function into batches. Remarkably, in our MCEM procedure, the fixed effects parameters can naturally be updated in mini-batches as the variance components usually hindering the decomposability are updated separately and conventional loss functions can be used. Recently, Simchoni and Rosset [2023] extended their framework to binary classification with one clustering feature by leveraging Gaussian quadrature to integrate over the random effects. However, unlike our approach, LMMNN is not applicable to classification with multiple clustering features and classes due to the inability of Gaussian quadrature to scale effectively to high-dimensional integrals [Bolker et al., 2009].

ARMED [Nguyen et al., 2023] is distinguished by the use of two additional networks besides the fixed effects network: one for predicting the effects of unseen clusters and one adversarial network to better disentangle clustering effects from the fixed effects. The model is trained using VI and evaluated solely for binary classification with a single random effect and cardinalities up to 20. Using high-cardinality features as random effects leads to a parameter explosion of the two additional networks, greatly increasing training time and memory requirements. Furthermore, the bias towards the prior distribution prevents estimating larger random effects, which leads to poor performance for high-variance scenarios. Our experiments in Subsection 3.1 have highlighted these limitations. Additionally, they have demonstrated that our approach scales effectively to high-dimensional random effects while offering accurate inter-cluster variance quantification.

Other existing approaches are limited to very small neural network architectures [Tandon et al., 2006; Mandel et al., 2021] or have been only implemented for regression tasks [Simchoni and Rosset, 2021; Avanzi et al., 2023; Kilian et al., 2023]. In addition, there are deep mixed effects models which use point estimates of random effects and thus are not comparable to the MENNs in our scope [Shi et al., 2022; W¨ortwein et al., 2023]. Furthermore, various approaches combining GLMMs and tree-based models were proposed [Sela and Simonoff, 2012; Hajjem et al., 2017; Ngufor et al., 2019; Sigrist, 2022].

As highlighted in Table 5, all existing approaches come with restrictions in applicability. Our approach is the first that applies to multi-class classification datasets with multiple random effects. It is worth noting that some of the existing approaches may in theory be extendable. However, none of the available implementations is readily applicable to the datasets we considered in Subsections 3.2 and 3.3 and the required modifications are nontrivial. Despite the discussed limitations, we want to emphasize that all approaches have particular strengths for specific scenarios. I.e., DeepGLMM can be applied to panel data, LMMNN is remarkably well suited for regression tasks and ARMED is the most efficient for low-dimensional random effects. The main limitation of all mixed-effects approaches is that they introduce components that increase the training time compared to conventional encoding or embedding approaches. MeNets and LMMNN rely on expensive matrix inversions, ARMED introduces two additional networks, and DeepGLMM as well as MC-GMENN rely on sampling the posterior. Nevertheless, we have shown that our approach is efficient compared to others in high-cardinality scenarios and has the widest range of applicability.

5 Conclusion

In this paper, we proposed MC-GMENN, a novel framework for training generalized mixed effects neural networks using Monte Carlo methods. We have shown that due to the partially Bayesian nature of MENNs, Monte Carlo methods can be utilized with competitive time efficiency. By decoupling batch updates from sampling in an MCEM procedure and state-of-the-art MCMC sampling techniques, random intercept neural networks for high-cardinality clustering features can be trained even more efficiently than using previous approaches. Furthermore, we demonstrated that our approach is able to correctly quantify inter-cluster variance, while previous approaches are biased and often unable to estimate the true posterior. Our contribution allows to apply mixed effects neural networks to a wide range of classification problems, including multiple classes and clustering features. Future work includes investigating whether MC-GMENN can improve over state-of-the-art approaches in specific domains, such as medicine, click-through rate prediction, or humancentered data applications. Moreover, we hope that our work inspires researchers to challenge assumptions about the scalability of Monte Carlo methods for deep learning applications in other fields than mixed effects modeling.

Acknowledgments

This research was supported in part by the German Federal Ministry for Economic Affairs and Climate Action of Germany (BMWK), and in part by the Ministry of Economic Affairs, Labour and Tourism Baden-W¨urttemberg.