UMEC: Unified model and embedding compression for efficient recommendation systems

1 INTRODUCTION . As the core component of a recommendation system ( RS ) , recommendation models ( RM ) based on ranking neural networks are widely adopted in general content recommendation and retrieval applications . In general , an effective recommendation model consists of two components : a feature embedding sub-model and a prediction sub-model , as illustrated in Figure 1 . Usually , an RM adopts neural networks to serve two sub-models . Formally , we denote an RM as f ( · ; W ) , whereW is the learnable parameters of f . For the inference , the model f takes the input feature data x to predict the confidence of the content , serving the recommendation applications . Specifically , we further define the embedding and prediction sub-models as fe ( · ; We ) and fp ( · ; Wp ) respectively , whereWe andWp are their own learnable parameters andW = { We , Wp } . The embedding feature , v : = fe ( x ; We ) , is the input of fp with the input data x . Hence , we can express the RM as f ( · ; W ) , fp ( fe ( · ; We ) ; Wp ) . Given a ranking training loss ` ( · ) ( i.e. , binary cross entropy ( BCE ) loss ) , the learning goal of the ranking model can be written as min W ∑ ( x , y ) ∈D ` ( f ( x ; W ) , y ) , where y is the ground-truth label , and D is the training dataset . Nowadays , extremely large-scale data have been poured into the recommendation system to predict user behavior in many applications . In the online inference procedure , the heaviest computation component is the layer-wise product between the hidden output vectors and the model parametersW , for a neural network-based RM . A slimmed neural network structure would save a great amount of power consumption during the inference . Hence , the main idea of an RM compression is to slim the structural complexity of f ( W ) and reduce the dimension of hidden output vectors . To obtain an efficient ranking model ( for example , MLP based ) for an RS , one may apply existing model compression methods to MLPs directly . For example , Li et al . ( 2016 ) removes the entire filters in the network together with their connecting feature maps in terms of magnitudes , which can also ∗Equal Contribution . be applied to MLP structures to remove a specific neuron as well as its connections . Molchanov et al . ( 2019 ) approximates the importance of a neuron ( filter ) to the final loss by using the first and second-order Taylor expansions , which can also be applied to pruning the neurons without hassle . There are also some compression methods focusing on dimension reduction of embedding feature vectors . Ginart et al . ( 2019 ) proposes a mixed dimension embedding scheme by designing non-uniform embedding dimensions scaling with popularity of features . Joglekar et al . ( 2020 ) uses Neuron Input Search ( NIS ) to learn the embedding dimensions for the sparse categorical features . ... ... ... Concatenation ... ...... Dense Input Sparse Input Embedding Sub-model Prediction Sub-model Embedding Feature ing neural network model and the prediction accuracy goal of the ranking model , with end-to-end gradient-based training . We reformulate the optimization of training loss associated with hard constraints into a minimax optimization problem solved by the alternating direction method of multipliers ( ADMM ) method ( Boyd et al. , 2011 ) . To summarize our contributions , we present the merits of UMEC as the following : • To the best of our knowledge , UMEC is the first unified optimization framework for the recommendation system scenario . Unlike those existing works that treat the selection of input feature and compression of the model as two individual problems , UMEC jointly learns both together via unifying both original prediction learning goal and the model compression related hard constraints . • We reformulate the joint input feature selection and model compression task as a constrained optimization problem . We convert resource constraints and L0 sparsity constraints into soft optimization energy terms and solve the whole optimization using ADMM methods . • Extensive experiments performed over large-scale public benchmarks show that our method largely outperforms previous state-of-the-art input feature selection methods and model compression methods , endorsing the benefits of the proposed end-to-end optimization . 2 RELATED WORK . Recommendation Models With the recent development of deep learning techniques , plenty of works have proposed learning-based recommendation systems to grapple with personalized recommendation tasks . Hidasi et al . ( 2015 ) applies recurrent neural networks ( RNN ) on long session data to model the whole session and achieve better accuracy than the traditional matrix factorization approaches . Covington et al . ( 2016 ) proposes a high-level recommendation algorithm for YouTube videos by specifying a candidate generation model followed by a separate ranking model , which demonstrates a significant performance boost thanks to neural networks . Wu et al . ( 2019 ) proposes a recommendation algorithm based on the graph neural network to better capture the complex transitions among items , achieving a more accurate item embedding . Their method outperforms traditional sequential approaches . Naumov et al . ( 2019 ) exploits categorical features and designs a parallelism scheme for the embedding tables to alleviate the limited memory problem . Huang et al . ( 2020 ) proposes a multi-attention based model to mitigate the low efficiency problem in the group recommendation scenarios . Input Feature Compression for Recommendation Models The numerous input features for practical recommendation scenarios necessitate the selection of useful ones to save memory resources and facilitate computational efficiency during inference . In factorization machines ( FM ) -based and context-aware recommendation systems , for example , Mao et al . ( 2017 ) conducts input feature selection in an automatic manner by feature ranking and feature sub-sampling , while managing to improve the prediction quality . Ginart et al . ( 2019 ) designs a mixed dimension embedding layers scheme and adjusts the dimension of a particular embedding feature according to the frequency of that item . Joglekar et al . ( 2020 ) performs Neural Input Search ( NIS ) method to search for the embedding sizes for the sparse categorical features , and designs the multi-size embedding framework to leverage the model capacity more efficiently . Song et al . ( 2020 ) proposes a discovering framework to automatically find the interaction architecture of the prediction model . Model Compression Model compression aims to reduce the complexity of Deep Neural Network ( DNN ) models to achieve inference efficiency and resource saving . Generally , model compression techniques can be categorized into pruning ( Han et al. , 2015 ; Wen et al. , 2016 ; Molchanov et al. , 2019 ) , quantization ( Hubara et al. , 2017 ; Wu et al. , 2018 ; Ajanthan et al. , 2019 ) and distillation ( Hinton et al. , 2015 ; Polino et al. , 2018 ; Tung & Mori , 2019 ) . Han et al . ( 2015 ) proposes a magnitude-based and element-wise pruning technique , which does not guarantee the reduction of computation efficiency , energy and memory costs . Structured pruning can solve this issue by removing whole certain types of structures from the network , such as filter pruning and channel pruning methods ( Li et al. , 2016 ; He et al. , 2017 ; Molchanov et al. , 2019 ) . Recent works ( Frankle & Carbin , 2019 ; Chen et al. , 2020b ; a ; 2021c ; b ; a ; Ma et al. , 2021 ) reveal that with appropriate compression techniques , the identified sparse subnetworks are capable of training in isolation to match the dense model performance . In the meanwhile , researchers also develop great interest in investigating hardware-aware compression , which utilizes practical resource requirements ( e.g. , energy , latency , Flops ) to guide the compression ( He et al. , 2018 ; Yang et al. , 2019 ; Liu et al. , 2019 ) . 3 UNIFIED MODEL AND EMBEDDING COMPRESSION ON RS . For an arbitrary NN-based recommendation model f ( · ) in an RS , the terminology of `` compression '' refers to reduce the power consumption with respect to all the computation operations for the inference . The learning parameters of f ( W ) can be expressed asW : = { W ( l ) |l ∈ [ L ] } , where L is the number of layers in the f ( W ) . For an arbitrary layer l , the layer ’ s weight is denoted as W ( l ) ∈ Rdout×din . The recommendation inference on the layer l is executed as x ( l+1 ) = σ ( W ( l ) x ( l ) + b ( l ) ) , where x ( l ) ∈ Rdin is the input of l-th layer , and b ( l ) is the corresponding learning bias term . Since the RM consists of two components : feature embedding sub-model and neural network prediction sub-model , the `` compression '' can be interpreted as shrinking the dimension of embedding feature v and reducing the number of parameters inW and b . Perspective of model parameters W and b It is to reduce the input or output neurons for the specific layer l. For example , the i-th neuron of x ( l+1 ) and j-th neuron of x ( l ) are chosen to remove , then the weight matrix W ( l ) will be reshaped to ( dout − 1 ) × ( din − 1 ) . Perspective of feature embedding vectors v The reduction of din would result in the elimination of the input features . For example , v ( 1 ) is defined as the concatenation of a series of embedding vectors , v ( 1 ) = [ e1 ; e2 ; · · · ; en ] , where [ · ; · ; · ] represents the concatenation operator , and n denotes the number of embedding feature vectors in the RS . Without loss of generality , for the first embedding vector e1 ∈ Rd1 , the zero-out of first d1 columns in W ( 1 ) will result in the elimination of the whole embedding vectors for e1 . On the other hand , the compression on fe ( We ) would reduce the feature embedding vectors ’ dimensions . 3.1 A COMPRESSION EXAMPLE : GROUP LASSO FOR WEIGHT PRUNING . We first consider a simple example with an RS model consisting of fe ( We ) and fp ( Wp ) . For this RS , to reduce the computation consumption , we present a straightforward model compression method with Group Lasso . The method considers both the feature embedding selection and network slimming . We only need to apply the sparse regularization terms to the weight matrices . The optimization problem can be expressed as min W ` ( W ) + λ1 ∑ i ∥∥∥W ( 1 ) . , gi∥∥∥ 1 + L∑ l=2 λl ∑ i ∥∥∥W ( l ) . , i ∥∥∥ 1 , ( 1 ) where ` ( W ) is the original objective function . The λs are hyper-parameters to control the importance of these regularization terms . g· denotes the feature groups for the input layer . ‖ · ‖1 indicates the L1 norm . The optimization problem ( 1 ) restricts the group-wise sparsity for the input layer and the neuron-wise sparsity ( or group size = 1 ) for the rest model layers while training with the original objective function ` ( W ) . The optimization problem can be solved with the projected gradient descent method . However , the optimization problem ( 1 ) has its own limitation . The hyper-parameters λs only have the latent impact on the model size and we could not directly affect the model size . It is not friendly when the users would like to restrict the model size to a specific number . On the other hand , the hyper-parameters also make the optimization problem hard to tune , requiring multiple arguments searching rounds . To avoid these shortages , we are supposed to have new approaches to handling the compression of RS .

Summary: The paper proposes a new unified optimization framework to solve the RM compression problem. It jointly compresses the prediction network and the feature embedding, and automatically optimizes their complexity budgets under the total resource constraint. In this way, it avoids treating selection of input feature and compression of model as two individual problems, and requires no expensive hyper-parameters tuning for either part. The novel unified model and embedding compression (UMEC) method can directly satisfy both the requirement of resource consumption of the ranking neural network model and the prediction accuracy goal of the ranking model, with end-to-end gradient-based training.