Learnable Embedding sizes for Recommender Systems

1 INTRODUCTION . The success of deep learning-based recommendation models ( Zhang et al. , 2019 ) demonstrates their advantage in learning feature representations , especially for the most widely-used categorical features . These models utilize the embedding technique to map these sparse categorical features into real-valued dense vectors to extract users ’ preferences and items ’ characteristics . The learned vectors are then fed into prediction models , such as the inner product in FM ( Rendle , 2010 ) , selfattention networks in AutoInt ( Song et al. , 2019 ) , to obtain the prediction results . The embedding table could contain a large number of parameters and cost huge amounts of memory since there are always a large number of raw features . Therefore , the embedding table takes the most storage cost . A good case in point is the YouTube Recommendation Systems ( Covington et al. , 2016 ) . It demands tens of millions of parameters for embeddings of the YouTube video IDs . Considering the increasing demand for instant recommendations in today ’ s service providers , the scale of embedding tables becomes the efficiency bottleneck of deep learning recommendation models . On the other hand , features with uniform embedding size may hard to handle the heterogeneity among different features . For example , some features are more sparse , and assigning too large embedding sizes is likely ∗Chen Gao is the Corresponding Author . The work is performed when Siyi Liu is an intern in Tsinghua University . 1Codes are available at : https : //github.com/ssui-liu/learnable-embed-sizes-for-RecSys to result in over-fitting issues . Consequently , recommendation models tend to be sub-optimal when embedding sizes are uniform for all features . The existing works towards this problem can be divided into two categories . Some works ( Zhang et al. , 2020 ; Shi et al. , 2020 ; Kang et al. , 2020 ) proposed that some closely-related features can share parts of embeddings , reducing the whole cost . Some other works ( Joglekar et al. , 2020 ; Zhao et al. , 2020b ; a ; Cheng et al. , 2020 ) proposed to assign embeddings with flexible sizes to different features relying on human-designed rules ( Ginart et al. , 2019 ) or neural architecture search ( Joglekar et al. , 2020 ; Zhao et al. , 2020b ; a ; Cheng et al. , 2020 ) . Despite a reduced embedding size table , these methods still can not perform well on the two most concerned aspects , recommendation performance and computation cost . Specifically , these methods either obtain poor recommendation performance or spend a lot of time and efforts in getting proper embedding sizes . In this paper , to address the limitations of existing works , we proposed a simple yet effective pruning-based framework , named Plug-in Embedding Pruning ( PEP ) , which can plug in various embedding-based recommendation models . Our method adopts a direct manner–pruning those unnecessary embedding parameters in one shot–to reduce parameter number . Specifically , we introduce the learnable threshold ( s ) that can be jointly trained with embedding parameters via gradient descent . Note that the threshold is utilized to determine the importance of each parameter automatically . Then the elements in the embedding vector that are smaller than the threshold will be pruned . Then the whole embedding table is pruned to make sure each feature has a suitable embedding size . That is , the embedding sizes are flexible . After getting the pruned embedding table , we retrain the recommendation model with the inspiration of the Lottery Ticket Hypothesis ( LTH ) ( Frankle & Carbin , 2018 ) , which demonstrates that a subnetwork can reach higher accuracy compared with the original network . Based on flexible embedding sizes and the LTH , our PEP can cuts down embedding parameters while maintaining and even boosting the model ’ s recommendation performance . Finally , while there is always a trade-off between recommendation performance and parameter number , our PEP can obtain multiple pruned embedding tables by running only once . In other words , our PEP can generate several memory-efficient embedding matrices once-for-all , which can well handle the various demands for performance or memory-efficiency in real-world applications . We conduct extensive experiments on three public benchmark datasets : Criteo , Avazu , and MovieLens-1M . The results demonstrate that our PEP can not only achieve the best performance compared with state-of-the-art baselines but also reduces 97 % to 99 % parameter usage . Further studies show that our PEP is quite computationally-efficient , requiring a few additional time for embedding-size learning . Furthermore , visualization and interpretability analysis on learned embedding confirm that our PEP can capture features ’ intrinsic properties , which provides insights for future researches . 2 RELATED WORK . Existing works try to reduce the embedding table size of recommendation models from two perspectives , embedding parameter sharing and embedding size selection . 2.1 EMBEDDING PARAMETER SHARING . The core idea of these methods is to make different features re-use embeddings via parameter sharing . Kang et al . ( 2020 ) proposed MGQE that retrieves embedding fragments from a small size of shared centroid embeddings and then generates final embedding by concatenating those fragments . Zhang et al . ( 2020 ) used the double-hash trick to make low-frequency features share a small embedding-table while reducing the likelihood of a hash collision . Shi et al . ( 2020 ) tried to yield a unique embedding vector for each feature category from a small embedding table by combining multiple smaller embedding ( called embedding fragments ) . The combination is usually through concatenation , add , or element-wise multiplication among embedding fragments . However , those methods suffer from two limitations . First , engineers are required to carefully design the parameter-sharing ratio to balance accuracy and memory costs . Second , these rough embeddingsharing strategies can not find the redundant parts in the embedding tables , and thus it always causes a drop in recommendation performance . In this work , our method automatically chooses suitable embedding usages by learning from data . Therefore , engineers can be free from massive efforts for designing sharing strategy , and the model performance can be boosted via removing redundant parameters and alleviating the over-fitting issue . 2.2 EMEBDDING SIZE SELECTION . The embedding-sharing methods assign uniform embedding sizes to every feature , which may still fail to deal with the heterogeneity among different features . Recently , several methods proposed a new paradigm of mixed-dimension embedding table . Specifically , different from assigning all features with uniformed embedding size , different features can have different embedding sizes . MDE ( Ginart et al. , 2019 ) proposed a human-defined rule that the embedding size of a feature is proportional to its popularity . However , this rule-based method is too rough and can not handle those important features with low frequency . Additionally , there are plenty of hyper-parameters in MDE requiring a lot of truning efforts . Some other works ( Joglekar et al. , 2020 ; Zhao et al. , 2020b ; a ; Cheng et al. , 2020 ) assigned adaptive embedding sizes to different features , relying on the advances in Neural Architecture Search ( NAS ) ( Elsken et al. , 2019 ) , a significant research direction of Automated Machine Learning ( AutoML ) ( Hutter et al. , 2019 ) . NIS ( Joglekar et al. , 2020 ) used a reinforcement learning-based algorithm to search embedding size from a candidate set predefined by human experts . A controller is adopted to generate the probability distribution of size for specific feature embeddings . This was further extended by DartsEmb ( Zhao et al. , 2020b ) by replacing the reinforcement learning searching algorithm with differentiable search ( Liu et al. , 2018 ) . AutoDim ( Zhao et al. , 2020a ) allocated different embedding sizes for different feature fields , rather than individual features , in a same way as DartsEmb . DNIS ( Cheng et al. , 2020 ) made the candidate embedding size to be continuous without predefined candidate dimensions . However , all these NAS-based methods require extremely high computation costs in the searching procedure . Even for methods that adopt differential architecture search algorithms , the searching cost is still not affordable . Moreover , these methods also require a great effort in designing proper search spaces . Different from these works , our pruning-based method can be trained quite efficiently and does not require any human efforts in determining the embedding-size candidates . 3 PROBLEM FORMULATION . Feature-based recommender system2 is commonly used in today ’ s information services . In general , deep learning recommendation models take various raw features , including users ’ profiles and items ’ attributes , as input and predict the probability that a user like an item . Specifically , models take the combination of user ’ s profiles and item ’ s attributes , denoted by x , as its ’ input vector , where x is the concatenation of all fields that could defined as follows : x = [ x1 ; x2 ; . . . ; xM ] , ( 1 ) where M denotes the number of total feature fields , and xi is the feature representation ( one-hot vector in usual ) of the i-th field . Then for xi , the embedding-based recommendation models generate corresponding embedding vector vi via following formulation : vi = Vixi , ( 2 ) 2It is also known as click-through rate prediction . where Vi ∈ Rni×d is an embedding matrix of i-th field , ni denotes the number of features in the i-th field , and d denotes the size of embedding vectors . The model ’ s embedding matrices V for all fields of features can be formulated as follows , V = { V1 , V2 , . . . , VM } , ( 3 ) The prediction score could be calculated with V and model ’ s other parameters ( mainly refer to the parameters in prediction model ) Θ as follows , ŷ = φ ( x|V , Θ ) , ( 4 ) where ŷ is the predicted probability and φ represent the prediction model , such as FM ( Rendle , 2010 ) or AutoInt ( Song et al. , 2019 ) . As for model training , to learn the models parameters , the optimizer minimizes the training loss as follows , min L ( V , Θ , D ) , ( 5 ) where D = { x , y } represents the data fed into the model , x denotes the input feature , y denotes the ground truth label , and L is the loss function . The Logloss is the most widely-used loss function in recommendation tasks ( Rendle , 2010 ; Guo et al. , 2017 ; Song et al. , 2019 ) and calculated as follows , L = − 1|D| |D|∑ j=1 ( yj log ( ŷj ) + ( 1− yj ) log ( 1− ŷj ) ) , ( 6 ) where |D| is the total number of training samples and regularization terms are omitted for simplification . 4 METHODOLOGY . 4.1 LEARNABLE EMBEDDING SIZES THROUGH PRUNING . As mentioned above , a feasible solution for memory-efficient embedding learning is to automatically assign different embedding sizes d̃i for different features embeddings vi , which is our goal . However , to learn d̃i directly is infeasible due to its discreteness and extremely-large optimization space . To address it , we propose a novel idea that enforce column-wise sparsity on V , which equivalently shrinks the embedding size . For example , as it shown in Figure 1 , the first value in embedding v1 is pruned and set to zero , leading to a d̃1 = d1 − 1 embedding size in effect . Furthermore , some unimportant feature embeddings , like v3 , are dropped by set all values to zero3 . Thus our method can significantly cut down embedding parameters . Note that the technique of sparse matrix storage help us to significantly save memory usage ( Virtanen et al. , 2020 ) . In such a way , we recast the problem of embedding-size selection into learning column-wise sparsity for the embedding matrix V. To achieve that , we design a sparsity constraint on V as follows , min L , s.t . ||V||0 ≤ k , ( 7 ) where || · ||0 denotes the L0-norm , i.e . the number of non-zeros and k is the parameter budget , which is , the constraint on the total number of embedding parameters . 3Our PEP benefit from such kind of reduction , as demonstrated in Section 5.1 , 5.3 and 5.4 . However , direct optimization of Equation ( 7 ) is NP-hard due to the non-convexity of the L0-norm constraint . To solve this problem , the convex relaxation of L0-norm , called L1-norm , has been studied for a long time ( Taheri & Vorobyov , 2011 ; Beck & Teboulle , 2009 ; Jain et al. , 2014 ) . For example , the Projected Gradient Descent ( PGD ) ( Jain et al. , 2014 ) in particular has been proposed to project parameters to L1 ball to make the gradient computable in almost closed form . Note that the L1 ball projection is also known as Soft Thresholding ( Kusupati et al. , 2020 ) . Nevertheless , such methods are still faced with two major issues . First , the process of projecting the optimization values onto L1 ball requires too much computation cost , especially when the recommendation model has millions of parameters . Second , the parameter budget k requires human experts to manually set at a global level . Considering that features have various importance for recommendation , such operation is obviously sub-optimal . To tackle those two challenges , inspired by Soft Threshold Reparameterization ( Kusupati et al. , 2020 ) , we directly optimize the projection of V and adaptively pruning the V via learnable threshold ( s ) which can be updated by gradient descent . The re-parameterization of V can be formulated as follows , V̂ = S ( V , s ) = sign ( V ) ReLU ( |V| − g ( s ) ) , ( 8 ) where V̂ ∈ RN×d denotes the re-parameterized embedding matrix , and g ( s ) serves as a pruning threshold value , of which sigmoid function is a simple yet effective solution.4 We set the initial value of trainable parameter s ∈ R ( called sinit ) to make sure that the threshold ( s ) g start close to zero . The sign ( · ) function converts positive input value to 1 and negative input value to -1 , and zero input will keep unchanged . As S ( V , s ) is applied to each element of V , and thus the optimization problem in Equation ( 5 ) could be redefined as follows , min L ( S ( V , s ) , Θ , D ) . ( 9 ) Then the trainable pruning parameter s could be jointly optimized with parameters of the recommendation models φ , through the standard back-propagation . Specifically , the gradient descent update equation for V at t-th step is formulated as follows , V ( t+1 ) ← V ( t ) − ηt∇S ( V , s ) L ( S ( V ( t ) , s ) , D ) ∇VS ( V , s ) , ( 10 ) where ηt is t-th step learning rate and denotes the Hadamard product . To solve the nondifferentiablilty of S ( · ) , we use sub-gradient to reformat the update equation as follows , V ( t+1 ) ← V ( t ) − ηt∇S ( V , s ) L ( S ( V ( t ) , s ) , D ) 1 { S ( V ( t ) , s ) 6= 0 } , ( 11 ) where 1 { · } denotes the indicator function . Then , as long as we choose a continuous function g in S ( · ) , then the loss functionL ( S ( V ( t ) , s ) , D ) would be continuous for s. Moreover , the sub-gradient of L with respect to s can be used of gradient descent on s as well . Thanks to the automatic differentiation framework like TensorFlow ( Abadi et al. , 2016 ) and PyTorch ( Paszke et al. , 2019 ) , we are free from above complex gradient computation process . Our PEP code can be found in Figure 7 of Appendix A.2 . As we can see , it is quite simple to incorporate with existing recommendation models , and there is no need for us to manually design the backpropagation process .

Summary: This paper proposed a novel approach to reduce size of the embedding table while not to drop in accuracy and computational optimization. Fixed-size embedding table has two problems, high memory usage cost and overfitting problem for those features that do not require too large representation. This paper recast the problem of embedding-size selection into learning column-wise sparsity, constraint K (eq(7)) and then convert S(V,s) problem (eq(8)). Paper used three benchmark datasets and some classical methods to verify effect.