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: The paper proposes PEP (Plug-in Embedding Pruning) to reduce the size of embedding table while incurring insignificant drop in accuracy. The related work is well summarized into Embedding Parameter Sharing and Embedding Size Selection methods and the motivation for the current approach is well explained. The paper draws inspiration from Lottery Ticket Hypothesis. The problem formulation of Embedding pruning is done in a crisp way avoiding additional hyper parameter tuning that can be found in other methods. Similar to LTH, the paper shows that the initiation strategy can make the training process faster and stable. The results show an impressive 97-99% parameter pruning via PEP. As for the computation cost, PEP results show an additional 20-30% time cost compare with base models.