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

1 Introduction . 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 . 2 A Brief Tour of Linear Mixed Models . 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 . 3 LMMNN : Our Proposed Approach . 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 )

Summary: 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.