Using Random Effects to Account for High-Cardinality Categorical Features and Repeated Measures in Deep Neural Networks
Year: 2021

Abstract:
High-cardinality categorical features are a major challenge for machine learning methods in general and for deep learning in particular. Existing solutions such as one-hot encoding and entity embeddings can be hard to scale when the cardinality is very high, require much space, are hard to interpret or may overfit the data. A special scenario of interest is that of repeated measures, where the categorical feature is the identity of the individual or object, and each object is measured several times, possibly under different conditions (values of the other features). We propose accounting for high-cardinality categorical features as random effects variables in a regression setting, and consequently adopt the corresponding negative log likelihood loss from the linear mixed models (LMM) statistical literature and integrate it in a deep learning framework. We test our model which we call LMMNN on simulated as well as real datasets with a single categorical feature with high cardinality, using various baseline neural networks architectures such as convolutional networks and LSTM, and various applications in e-commerce, healthcare and computer vision. Our results show that treating high-cardinality categorical features as random effects leads to a significant improvement in prediction performance compared to state of the art alternatives. Potential extensions such as accounting for multiple categorical features and classification settings are discussed. Our code and simulations are available at https://github.com/gsimchoni/lmmnn

heading: 1 Introduction
text: In recent years deep neural networks (DNNs) have served as the method of choice for learning complex non-linear relations between features of large datasets and for hard prediction tasks. Yet, despite their advanced machinery such as stochastic gradient descent (SGD) and convolutional layers, DNNs (as other machine learning tools) do not readily lend themselves to handling categorical features of high cardinality. Such features are often seen in modeling tasks, a good recent example would be Lin et al. [16] who used electronic medical records (EMR) of hospital patients to predict hospital readmission, and had to deal with the patient’s diagnosis as a categorical feature with thousands of levels. In a recent review Hancock and Khoshgoftaar [13] surveyed the approaches taken by DNN practitioners for turning the levels of categorical features into numeric form, of which the most used approaches are one-hot encoding (OHE) and entity embeddings. OHE and embeddings also appear as the go-to solution in one of the most cited handbooks of machine learning by Géron [9]. See Section 4 for a more detailed discussion of these approaches.
A different approach for handling high-cardinality categorical features is that of linear mixed models (LMM), see e.g. McCulloch et al. [18]. In the LMM framework statisticians differentiate between
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
fixed effects (FE) which are regular parameters to be estimated from the data and random effects (RE), which are parameters that are not directly estimated, but rather treated as random variables, and their distribution is estimated instead. This mechanism is often used to deal with high cardinality categorical variables or "cluster" indicators. In a typical setting a statistician would be interested in the effect few designed experimental conditions have on a dependent variable, e.g. the effect of treatment A versus treatment B on a patient’s blood pressure. The treatment effect would be considered a fixed effect, while the effect of the city in which a subject resides — a categorical feature with possibly thousands of levels — would be considered a RE variable, often having normal distribution N ( 0, σ2b ) with σ2b being an unknown "variance component" to be estimated (see Section 2 for a formal description). Another common example of a RE variable is the subject of the experiment herself, in a repeated measures setting, where a subject is treated to both treatments A and B. In such a setting a statistician would usually deem a subject’s effect as random, i.e. coming from a possibly infinite population of subjects, and use proper LMM machinery to handle such a variable.
In this paper we are interested in examining modern big-data predictive modeling domains, where DNNs would be a natural choice of modeling approach, but where high dimensional categorical variables or repeated measures are an important aspect of prediction. Examples we pursue might include: predicting the gaze direction of a person’s eyes from repeated facial images of the same individuals [24], detection of diabetic retinopathy from retinal fundus photographs where each patient is photographed for both left and right eyes [11], prediction of a person’s fluid intelligence from dozens of features including a person’s job which can take hundreds of levels [14] and predicting the price of an Airbnb rental with host being one of tens of thousands [15].
To address such challenges we seek a simple extension to DNN which allows to treat high-cardinality categorical features as RE and ultimately improve the network’s prediction performance in a regression setting. We do this by generalizing the LMM-framework negative log likelihood (NLL) function and treating it as a natural loss function for the entire network, minimizing both the "fixed" and "random" parts, and improving the bottom line test mean squared error (MSE). We demonstrate the superiority of our method over using existing methods such as OHE, entity embeddings and standard LMM via R’s lme4 package (Bates et al. [2]), as well as previous attempts at integrating RE in DNN (MeNets, Xiong et al. [22]), on both simulated and real datasets from various applications such as healthcare and computer vision. Finally we discuss future directions, such as using RE in DNN in a classification setting and with more complex covariance structures such as kriging.

heading: 2 A Brief Tour of Linear Mixed Models
text: The canonical linear model assumes: y = Xβ + , (1) where X is the n×p model matrix where the ith row is xi, β ∈ Rp is a vector of model parameters to be estimated and ∈ Rn is normal i.i.d noise, i.e. ∼ N ( 0, σ2eI ) , where σ2e is a variance parameter and I is the n× n identity matrix. A linear mixed model treats the β parameters as fixed effects, and allows additional data to enter the model in the form of a n× q matrix Z and a vector of random effects b ∈ Rq:
y = Xβ + Zb+ (2) Here b are random variables, typically assumed to have a multivariate normal distribution N (0, D) where D is a q × q positive semi-definite matrix holding usually unknown variance components to be estimated, let these be θ, so D could be written as D(θ). The structure of this covariance matrix is up to the researcher but there are typically simplified structures used. It is further assumed that there is no dependence between the normal noise and the random effects, i.e. cov ( , b) = 0.
Let us write the vector of all variance components as ψ = [σ2e , θ]. How to estimate β, ψ? The statistical approach of choice is maximum likelihood estimation (MLE). We write the marginal distribution of y as N (Xβ, V ) where V (ψ) = ZD(θ)Z ′ + σ2eI , and from here it is straightforward to see that the log likelihood is:
l(β, ψ|y) = −1 2 (y −Xβ)′ V (ψ)−1 (y −Xβ)− 1 2 log |V (ψ)| − n 2 log 2π (3)
The MLEs for β, ψ maximize l(β, ψ|y) or equivalently minimize the negative log likelihood (NLL) in (3). Typically maximizing the likelihood is difficult with given constraints on ψ (which are nonnegative variance components), and to get less biased estimates for ψ alternatives methods such as
restricted maximum likelihood estimation (REML) are preferred. For full details see e.g. McCulloch et al. [18].
Of particular interest is the random intercepts model, in which a single categorical random variable, such as "subject" or "city", is assumed. If this so-called "clustering" variable has q levels, then Z is an indicator matrix holding at entry [i, j] 1 if observation i belongs to cluster j and zero otherwise. In this case it is often assumed that the random effects for this variable are independent of one another, hence: D = σ2b I . In this setting V (ψ) = σ 2 bZZ
′ + σ2eI is a block diagonal matrix and its inversion can be avoided in (3). More on that in Section 3.
How to predict ŷte in a machine learning scenario, where (X,Z, y) are typically split into training and testing sets (Xtr, Ztr, ytr) and (Xte, Xte, yte)? In general one would use y’s conditional distribution:
ŷte = Xteβ̂ + Zteb̂, (4)
where β̂ = (X ′trV̂ −1Xtr) −1X ′trV̂ −1ytr are the estimated fixed effects once the estimated variance components ψ̂ are input into V̂ , and:
b̂ = D̂Z ′trV̂ −1 ( ytr −Xtrβ̂ ) (5)
is the so called best linear unbiased predictor (BLUP), as b are not actually parameters to be estimated, but random variables to be predicted.
Sometimes, however, (4) is not possible such as in the case of the random intercepts model, where Zte holds levels of the random variable unseen before, e.g. a new subject. In this case it is customary to use y’s marginal distribution and predict ŷte to be Xteβ̂, i.e. without the random part. Finally, still under the random intercepts model, it can be shown that for a given level j and corresponding random effect b̂j , the computation can be simplified to avoid the inversion of V̂ giving:
b̂j = nj σ̂
2 b
σ̂2e + nj σ̂ 2 b
( ȳtr;j −Xtrβj ) , (6)
where ( σ̂2e , σ̂ 2 b ) are the estimated variance components, nj is the no. of observations in cluster j and ȳtr;j and Xtrβj are the true and predicted mean values of y in cluster j respectively.

heading: 3 LMMNN: Our Proposed Approach
text: Let us change (2) into: y = f (X) + g (Z) b+ , (7)
where f and g are non-trivial functions which DNNs are suitable for, e.g. non-linear and involving interactions. We propose to use NLL as a natural loss function, where Xβ is replaced by the DNN outputs f(X):
NLL(f, g, ψ|y) = 1 2 (y − f (X))′ V (g, ψ)−1 (y − f (X)) + 1 2 log |V (g, ψ)|+ n 2 log 2π, (8)
where V (g, ψ) = g(Z)D(θ)g(Z)′+σ2eI . Note that f and g are kept general as possible, to allow any acceptable DNN architecture, including convolutional and recurrent neural networks. See Figure 1 for a schematic description of our approach which we call LMMNN, in the case f and g are approximated with a simple multilayer perceptron.
We propose using existing DNN machinery, mainly SGD, to fit the model and its variance components. To be more specific, we implemented a custom NLL loss layer where at each epoch (8) is calculated on a small batch typically of size 30-50 observations and auto-differentiation is handled by Keras [4]. An alternative could be using explicit formulas for the variance components ψ derivatives in case V (ψ) is simple:
∂NLL ∂ψ = −1 2 (y − f (X))′ V −1 ∂V ∂ψ V −1 (y − f (X)) + 1 2 tr
( V −1 ∂V
∂ψ
) , (9)
where ∂V∂ψ is g(Z) ∂D(θ) ∂θ g(Z) ′ for the θ variance components and I for σ2e .
The main challenge in implementing mini-batch SGD on (8) is the inverse and determinant of V (g, ψ) which may not decompose into sums on batches. An important exception is the main setting of interest in our paper in the random intercepts model, when g(Z) = Z and V (g, ψ) is a block-diagonal matrix. If all blocks are identical, such that each level of the categorical feature has m observations (Xj , yj) where j = 1, . . . q, we can write V (ψ) = diag(V1, ..., Vq) where each Vj block is of size m×m and V (ψ)j = σ2bJm+σ2eIm where Jm is am×m all 1s matrix. This means we can write the inverse in (8) as block diagonal as well, V (ψ)−1 = diag(V −11 , ..., V −1 q ), and the log determinant in
(8) as a sum of log determinants: log |V (ψ)| = ∑q j=1 log |Vj |. The NLL in (8) can now be written as
a sum: NLL(f, ψ|y) = ∑q j=1 1 2 (yj − f (Xj)) ′ V −1j (yj − f (Xj)) + 12 log |Vj |+ m 2 log 2π. Most importantly, the gradient in (9) can be decomposed into a sum of gradients over mini-batches of size m:
∂NLL
∂ψ = q∑ j=1 [ −1 2 (yj − f (Xj))′ V −1j ∂Vj ∂ψ V −1j (yj − f (Xj)) + 1 2 tr ( V −1j ∂Vj ∂ψ )] (10)
If all blocks are not identical but small (the typical case), we can still make sure that mini-batches consist of each level j’s nj observations. For more general cases the mini-batch approach (where inverse and determinant are calculated on the mini-batch) still seems to work well in practice, some related theory can be found in Chen et al. [3], see Section 6 for further discussion on more LMM scenarios.
For prediction, we accommodate (4) as well:
ŷte = f̂ (Xte) + ĝ (Zte) b̂, (11)
where f̂ and ĝ are the outputs of the DNNs used to approximate f and g, and b̂ are the predicted random effects with the modified BLUP: b̂ = D(ψ̂)ĝ (Ztr) ′ V (ĝ, ψ̂)−1 ( ytr − f̂ (Xtr) ) Note further that in the likely case where g is simply the identity matrix and we are interested in a single clustering variable as in the random intercepts model, we can also avoid inversion of V in the loss function as described in [18], as well as in computing the modified BLUP:
b̂j = nj σ̂
2 b
σ̂2e + nj σ̂ 2 b
( ȳtr;j − f̂ (Xtr)j ) (12)

heading: 4 Related Work

heading: 4.1 One-Hot Encoding and and Entity Embeddings
text: As Hancock and Khoshgoftaar [13] state, OHE is often a "first step", and the standard method for handling categorical features. If v of length n is a categorical feature of q distinct levels, OHE would add q binary features z1, ..., zq , one for each level, with zli = 1 if observation i has level l in feature v, and 0 otherwise. The main advantages of the OHE approach are that it is deterministic, fast and highly explainable. On the downside it is hard to scale — as q can be in the tens or hundreds of thousands as in the examples below — a burden to any algorithm even when accounting for sparse representation of features. A second disadvantage of OHE features is how little information each feature carries. Other methods for converting categorical features into numeric can reach a much more meaningful representation.
One such method is entity embeddings (See e.g. Guo and Berkhahn [12]). Entity embeddings start by feeding a neural network with feature v after it has been one-hot encoded into a q-length vector, and using the network’s inherent loss function and back propagation algorithm to learn a low-dimensional representation of length d (typically d << q) to each of v’s q levels. This results in a lookup table E of dimensions q × d which serves the predictive modeling task and can also later be re-used via transfer learning where the representation learned for one task can serve for other tasks, see e.g. Do and Ng [6]. It is also our experience that the entity embeddings approach performs well. Apart from that, efficient implementations make sure it scales well when q is large, unlike OHE the resulting
representations may carry a lot of information, and the option to perform transfer learning on these is appealing. On the other hand, entity embeddings may consume much space (the E lookup table), they may need to be learned for each new task and the resulting representations are usually hard to interpret.

heading: 4.2 Previous Attempts at Combining Random Effects in DNN
text: There have been few previous attempts at combining random effects in DNN. The efforts we are aware of are MeNets by Xiong et al. [22] and DeepGLMM by Tran et al. [21].
MeNets Inspired by a non-linear mixed effects model y = ν (Xβ + Zb) + , where ν is some non-linear function, Xiong et al. [22] and Xiong et al. [23] propose the following model to learn fixed effects β and random effects b:
y = f (X)β + f (X) b+ (13)
There are two main differences between (13) and our (7): First, the Z RE features matrix is missing, rather the RE features are learned as the output of the neural network to a standard input X , i.e f(X). Second, (13) uses a single neural network to learn both the fixed and random features, where (7) is more general and allows different functions f, g for fixed and random features respectively. In particular, (7) allows for g to be the identity function which may certainly be appropriate for high cardinality categorical features and repeated measures and gives intuitive meaning to b.
In order to learn β and b the authors use variational expectation maximization (V-EM) combined with SGD: A E-Step in which β̂, b̂ are updated while minimizing the standard squared loss with a DNN, followed by a M-Step where the variance components ψ̂ are updated so as to maximize a NLL loss similar to (8). Hence, an additional critique of MeNets could be that they essentially use two loss functions for a single task, whereas we use NLL as a single loss function to be maximized in the SGD framework, making our implementation simpler and easier to track. In addition, in order to update b̂j at each E-step, there is the need to invert V̂j which is the jth RE covariance matrix of size nj × nj . This can be hard to scale if one of the categorical feature levels has a large nj 1
1See Section 4 and the UKB dataset, when one of the levels of the categorical feature "job" has over 29,000 observations - being unemployed - and MeNets inverts a 29,000 × 29,000 matrix at each epoch.
Finally we note that MeNets were demonstrated in a limited context: the authors put special emphasis on computer vision applications, mainly that of gaze estimation, as a consequence the DNN architecture used by the authors is solely a convolutional network, whereas we demonstrate our results on a broad class of DNN architectures. Moreover, the categorical clustering feature MeNets used was the subject whose gaze was recorded, with typically only 10-50 subjects (and many repeated measures for each subject). Below we demonstrate MeNets in a wider context and compare it to LMMNN, illustrating the superiority of our approach in terms of both predictive performance and computation.
DeepGLMM A different approach taken by Tran et al. [21] is based on a very specific mixed effects model, in which each subject i is repeatedly measured at different times t for some response yit which can be continuous as well as discrete, as modeled by generalized linear models (GLM). In such a model it makes sense to not only have a random intercept for each subject but also a random slope ai. The authors write:
g(µit) = f(x (1) it , w, β (1)) + (β(2) + ai) ′x (2) it , (14)
where w are the network parameters, x(1) and x(2) are the features expected to have nonlinear and linear effects, respectively, β(1) and β(2) are the fixed nonlinear and linear effects respectively, and µ = E(y|x), via some link function g, e.g. the logit function for binary y. A Bayesian approach based on variational approximation is used to maximize the likelihood for (14).
We note that (14) is similar to our criterion in (7), when g is the identity function and y is linear in Z, the RE features matrix. We further acknowledge that our approach is only currently applicable for a continuous y (though see future directions in Section 6). However, the variational approximation algorithm proposed in DeepGLMM, which combines numerous elements such as importance sampling, factor covariance, variable selection as well as choice of priors makes it challenging to implement, let alone use as a "plug-in" for different DNN architectures as we strive to do. Due to the high conceptual and computational complexity of DeepGLMM we were not able to implement and test it below like MeNets — it cannot scale to problems of the size we consider and to diverse architectures.

heading: 5 Results

heading: 5.1 Simulated Data
text: We start with a simple random intercepts model simulation, in which n = 100, 000 and there is a single categorical random variable with q levels (q ∈ {100, 1000, 10000}), such that Z is a n × q binary matrix, and b is a q-length vector of i.i.d random effects, sampled from a N ( 0, σ2b ) distribution (σ2b ∈ {0.1, 1, 10}). The q levels are not evenly distributed among the N observations, rather we use a multinomial distribution sampling, for more details see our code and description in additional material. There are 10 fixed features in X non-linearly related to y:
y = (X1 + · · ·+X10) · cos(X1 + · · ·+X10) + 2 ·X1 ·X2 + Zb+ (15)
Notice in this case g is assumed to be the identity function. We sample X features from a U (−1, 1) distribution, from a N ( 0, σ2e ) distribution where we keep σ2e = 1 always. We perform 5 iterations for each (q, σ2b ) combination (9 combinations in total), in which we sample the data, randomly split it into training (80%) and testing (20%), train our models to predict ŷte and compare the bottomline MSEs in predicting yte. We compare our model’s MSE to those of R’s lme4 package results (i.e. standard LMM), MeNets, OHE, entity embeddings and ignoring the categorical feature in Z altogether. We use the same DNN architecture for all neural networks, that is 4 hidden layers with 100, 50, 25, 12 neurons, a ReLU activation and a Dropout of 25% in each, and a final output layer with a single neuron. The loss we use is squared error loss (MSE) for OHE, embeddings and ignoring the RE, and NLL for LMMNN and MeNets (as mentioned above, MeNets uses squared loss for estimating fixed effects and NLL for variance components only). For both losses we use a batch size of 30 and an early stopping rule where training is stopped if no improvement in validation loss is seen within 10 epochs. We initialize both σ̂2e , σ̂ 2 b to be 1.0 where appropriate: R’s lme4 and LMMNN, and compare the resulting final estimates for these two methods. All runs are made on a Nvidia Quadro P620 GPU on a Windows machine, with Keras [4] and Tensorflow [1], with our own NLL loss layer for LMMNN, available on Github.
Table 1 summarizes the test MSE results and Table 2 (left) summarizes the estimated variance components results. As can be seen LMMNN reaches the smallest test MSE on average and with a considerable gap from the other methods, when standard errors are taken into account. This is particularly true when RE variance σ2b and cardinality q are high. As for the estimated variance components σ̂2e , σ̂ 2 b , LMMNN reaches a good estimation for both, while R’s lme4 reaches a poor estimation for σ2e without adding appropriate non-linear and interaction terms, resulting in worse prediction performance. We also plot in additional material Figure 1 the predicted random effects b̂ vs. the true random effects b for different qs, the category size distribution and also ŷte vs. yte. Finally Table 1 in additional material summarizes mean runtime and number of epochs.
A more challenging scenario is when g is some linear transformation Wq×d of Z into a lower dimension d:
y = (X1 + · · ·+X10) · cos(X1 + · · ·+X10) + 2 ·X1 ·X2 + ZWb+ , (16) where we use d = 0.1 · q and sample W from a U (−1, 1) distribution. b is now a d-length vector of random effects, sampled from a N ( 0, σ2b ) distribution (σ2b ∈ {0.1, 1, 10}). We use a Keras Embedding layer on Z to learn W , and notice now we cannot avoid inversion of V when calculating the loss in (8) or in predicting b̂ in (5). V is a N ×N matrix, so it is perfectly reasonable to invert it for a batch size of 30, but for predicting b̂ on the entire training set at the end of training this means creating and inverting a matrix of dimensions 80,000 × 80,000. In practice we sample 10,000 observations from the training set for estimating V and this seems to work well, as summarized by Table 3 which shows the test MSE scores for all methods. Our method performs the best especially at high cardinality q where the reduction in MSE can go over 80%. We do note that LMMNN struggles
in finding good σ̂2e , σ̂ 2 b estimates in such a complex scenario, as summarized in Table 2 (right), but its estimates are still substantially better than those of R’s lme4.

heading: 5.2 Real Data
text: We first describe the datasets and prediction tasks used in this work.
Estimating physical activity from self-reported behaviors from the UK Biobank The UK Biobank is an ongoing large scale cohort study where over 500,000 individuals across the UK aged 40-69 were first invited to 22 assessment centers in 2006-2010, surveyed and measured for various behaviors and metrics, such as gender, job, personal habits, current and historical health status [20]. Participants’ information has been de-identified and strict protocols exist for obtaining and disposing of the UK Biobank data. Several sub-samples of this cohort have also been invited to participate in more expensive and time consuming surveys such as neuroimaging and genome sequencing. One sub-sample of over 100,000 participants wore a wrist-mounted accelerometer for a period of 7 days [7]. We follow Pearce et al. [19] who attempted to predict these accelerometer-based physical activity (PA) data from much more readily available 14 self-reported behaviors such as no. of hours watching TV and sleeping, and gender. The motivation behind doing so is that physical activity is a strong marker for diseases outcomes such as cancer and respiratory disease mortality, yet it is hard to measure. The authors used linear regression and ignored some important features. Of particular interest here is the current job feature, a categorical feature with over 300 levels, unevenly distributed across participants. We attempt to improve prediction of physical activity by adding the job feature and using DNN in a LMMNN approach.
Estimating Drugs 1-10 rating from textual reviews from Drugs.com The website drugs.com is an online pharmaceutical database of drugs, providing information, reviews and ratings of thousands of doctor prescribed drugs. Gräßer et al. [10] scraped drugs.com for over 215,000 anonymized reviews and ratings which are available in the UCI Machine Learning Repository [8]. The authors discretized the 1-10 ratings into "negative", "neutral" and "positive" sentiments and used logistic regression on the text reviews processed to n-grams, to predict sentiment in a classification framework. Here we predict the rating itself in a regression framework using DNN, with a standard word embeddings layer followed by a LSTM layer and a single neuron output. A high-cardinality categorical feature which should improve prediction is the drug itself and there are over 3,600 drugs reviewed and rated in this dataset (so each drug has about 60 reviews on average, but these are unevenly distributed).
Landmark localization in CelebA facial images The CelebA dataset (Liu et al. [17]) contains 202,599 cropped facial images from 10,177 celebrities, where each celebrity has between 1 and 35 images, annotated for various attributes and landmarks such as the location of the tip of the nose and mouth. We use two convolutional neural networks to predict the (X,Y ) location of the nose and expect that treating facial images from the same individual as repeated measures should improve prediction. Therefore, the celebrity identity itself is the RE feature used here in the LMMNN framework.
Predicting prices of Airbnb rentals The website insideairbnb.com scrapes publicly available data from Airbnb rentals under the CC0 1.0 license [5]. Kalehbasti et al. [15] attempted to predict the log price of a rental from various features such as neighborhood, number of rooms and textual features extracted from reviews, overall 755 features. After performing variable selection they used multiple methods including DNN and found support vector regression (SVR) to give the best result in terms of test MSE. Their dataset holds about 50,000 listings with over 39,000 hosts which they did not even consider when performing variable selection. Some hosts publish only 1 rental, while some over 100.
In each dataset we use the same baseline DNN architecture detailed in the original analysis, if available. If the original analysis did not include the use of DNNs or did not provide specific details, we use a standard architecture such as the LSTM used for modeling the Drugs dataset reviews. The RE variable is either ignored, goes through embeddings or OHE (if possible computationally) before entering the network, used in the MeNets framework or used in the LMMNN framework, where we use g(Z) = Z always. For specifics on baseline network architecture for each dataset see additional material. We perform 5-fold cross validation where in each iteration we use 10% of the train set (which is 80% of the entire data) as validation data, and train the network until no improvement in validation loss is seen within 10 epochs. All runs are made on Google Colab Pro, with a Nvidia Tesla P100 GPU.
Table 4 describes the datasets in general, Table 5 summarizes the test MSE results and Table 3 in additional material summarizes mean runtime and number of epochs. It can be seen that including the high-cardinality variable in each of these datasets within a LMMNN framework achieves the lowest test MSE. We also refer the reader to additional material Figure 2 where we show true test y vs. predicted y scatter plots. Comparing to the results of the original papers is challenging since not all have supplied data and methods in a fully reproducible way and they use different metrics. Notwithstanding, we can verify our results are at least comparable in performance: For estimating physical activity in the UK Biobank, Pearce et al. [19] report achieving a test R2 of between 14 and 17%, while LMMNN achieves on average 19% (higher is better). For estimating drugs rating from text reviews we cannot compare to Gräßer et al. [10] results as they worked in a classification setting. For localizing the nose in the CelebA facial images, we are not aware of any MSE benchmarks but note that the LMMNN result of less than 2 pixels on average seems satisfactory. Finally, for the Airbnb rental price prediction, Kalehbasti et al. [15] report achieving a test MSE measure of 0.157 for a neural network, while LMMNN achieves on average a mean test MSE of 0.142, which is even better than the best result they achieved with SVR, 0.147.

heading: 6 LMM Extensions
text: The current paper is focused on regression settings with a single categorical feature, e.g. the subject with repeated measurements, as we wanted to take advantage of the computationally attractive LMM
feature described in equation (6) of avoiding any matrix inversions in calculating the predicted RE and the NLL loss. However, in further experimenting with LMMNN we note that it seems to perform well with more complex mixed effects scenarios which do necessitate inverting part of the V covariance matrix in equation (8) on each mini-batch during training. Apart from the simulation in which we defined g(Z) = ZW and used embeddings in the network architecture, other scenarios include having multiple categorical features, the random slopes model which adds a "slope" random effect to each categorical level, and even kriging over random fields, where V is assumed to contain spatio-temporal dependencies.
To give the reader a flavour for how LMMNN can be generalized to more scenarios, consider a situation of two uncorrelated categorical features, having q1 and q2 levels. Now Z in (7) is an indicator matrix of size n× (q1 + q2), b of size q1 + q2 distributes as before N (0, D), only now D is of size (q1 + q2)× (q1 + q2), and has structure:
D = [ σ2b1Iq1 0q1×q2 0q2×q1 σ 2 b2Iq2 ] , (17)
Now ψ holds three variance components [σ2e , σ 2 b1, σ 2 b2], and we can write the NLL and its derivatives in (8) and (9) exactly as before. However, notice that V in this case is not necessarily block-diagonal anymore, and LMMNN calculates the inverse on part of it for each mini-batch.
We show in Table 6 the mean test MSE results of a simulation similar to the first simulation in Section 5.1, where g is the identity matrix, f is in (15), σ2e = 1, n = 100, 000, using the same DNN architectures and rules. In fact the only difference from the simulation in Section 5.1 is that now we have two high-cardinality uncorrelated categorical features, we vary q1, q2, σ2b1, σ 2 b2 and LMMNN now has to estimate three variance components. As can be seen, LMMNN’s performance is superior to all other methods. It also estimates the variance components quite accurately when compared to R’s lme4, with a certain cost in runtime, see Tables 4 and 5 in additional material.

heading: 7 Conclusion
text: In this work we propose to treat high-cardinality categorical features in DNN as RE variables, to avoid the possible pitfalls of overfitting, over-parameterization and scalability issues often encountered when opting to use OHE or entity embeddings. LMMNN introduces a novel negative log likelihood loss function inspired by the well researched linear mixed effects model and its positive effect on DNN prediction performance is shown, with a simple implementation in Keras which can be plugged in to almost any regression DNN architecture, and is freely available on Github.
Treating categorical features with many levels in DNN as random effects seems to be promising as our simulated and real data results show and we intend to further pursue this direction in future work.
Acknowledgments and Disclosure of Funding
This study was supported in part by a fellowship from the Edmond J. Safra Center for Bioinformatics at Tel-Aviv University, and by Israeli Science Foundation grant 1804/16. UK Biobank research has been conducted using the UK Biobank Resource under Application Number 56885.

Review summaries:
This paper proposes to treat high-cardinality categorical features as the random effects of a mixed non-linear model. The proposed approach, called LMMNN, consists of minimizing the negative log-likelihood over both the functional estimates (parameterized as neural networks) and the variance components of the mixed non-linear model. The framework can be seen as an extension of the random intercepts model. The pipeline is eligible for automatic differentiation, though a closed form can be deduced in simple settings. Experiments are performed on both synthetic and various real datasets.

This paper proposes a random effects based objective for deep neural networks. The idea is to modify the standard linear mixed effects objective by replacing the linear predictor with the output of a neural network. The resulting objective can be trained with regular SGD. By modelling a categorical variable (e.g., subject identity) as a random effect, the result is argued to better account for this variable’s influence on the target variable compared to one-hot encoding, and embedding approaches. Experiments on synthetic and real-world datasets demonstrate the feasibility of the proposed approach.

The paper presented a methodology for dealing with high-cardinality features extracted from categorical variables by using the one-hot encoding applied to deep learning models. The suggested methodology decomposed the model into two sections: fixed effects (FE) and random effects (RE) and it is based on the negative log-likelihood loss from the linear mixed models for neural networks. Presented results including simulated data and four different real data sets with several dozens of categorical variables.