==== Front Front Big Data Front Big Data Front. Big Data Frontiers in Big Data 2624-909X Frontiers Media S.A. 10.3389/fdata.2023.1195742 Big Data Original Research PME: pruning-based multi-size embedding for recommender systems Liu Zirui 1 Song Qingquan 2 Li Li 3 Choi Soo-Hyun 3 Chen Rui 3 Hu Xia 1 * 1Computer Science Department, Rice University, Houston, TX, United States 2Linkedin, Sunnyvale, CA, United States 3Samsung Electronics America, Mountain View, CA, United States Edited by: Bo Han, Hong Kong Baptist University, Hong Kong SAR, China Reviewed by: Quanming Yao, Tsinghua University, China; Jiangchao Yao, Shanghai Jiao Tong University, China; Zhanke Zhou, Hong Kong Baptist University, Hong Kong SAR, China *Correspondence: Xia Hu xia.hu@rice.edu 15 6 2023 2023 6 119574228 3 2023 25 4 2023 Copyright © 2023 Liu, Song, Li, Choi, Chen and Hu. 2023 Liu, Song, Li, Choi, Chen and Hu https://creativecommons.org/licenses/by/4.0/ This is an open-access article distributed under the terms of the Creative Commons Attribution License (CC BY). The use, distribution or reproduction in other forums is permitted, provided the original author(s) and the copyright owner(s) are credited and that the original publication in this journal is cited, in accordance with accepted academic practice. No use, distribution or reproduction is permitted which does not comply with these terms. Embedding is widely used in recommendation models to learn feature representations. However, the traditional embedding technique that assigns a fixed size to all categorical features may be suboptimal due to the following reasons. In recommendation domain, the majority of categorical features' embeddings can be trained with less capacity without impacting model performance, thereby storing embeddings with equal length may incur unnecessary memory usage. Existing work that tries to allocate customized sizes for each feature usually either simply scales the embedding size with feature's popularity or formulates this size allocation problem as an architecture selection problem. Unfortunately, most of these methods either have large performance drop or incur significant extra time cost for searching proper embedding sizes. In this article, instead of formulating the size allocation problem as an architecture selection problem, we approach the problem from a pruning perspective and propose Pruning-based Multi-size Embedding (PME) framework. During the search phase, we prune the dimensions that have the least impact on model performance in the embedding to reduce its capacity. Then, we show that the customized size of each token can be obtained by transferring the capacity of its pruned embedding with significant less search cost. Experimental results validate that PME can efficiently find proper sizes and hence achieve strong performance while significantly reducing the number of parameters in the embedding layer. neural network recommender system embedding compression pruning scalability This work was funded by NSF IIS-2224843 and IIS-1849085. section-at-acceptanceMachine Learning and Artificial Intelligence ==== Body pmc1. Introduction Embedding feature information into vector representations is crucial for the success of deep learning based recommendation models (Zhang et al., 2019). In practice, the input features to recommender systems are often categorical, such as userID, itemID, and the category of items. For deep learning based recommendation models, these categorical features are mapped to low-dimensional learnable vectors (i.e., embeddings). Then, the learned vectors are fed into the rest of the model to learn the interaction between features. The number of layers in the rest of the recommendation model is typically small (usually less than 10) and independent of the number of categorical features (Cheng et al., 2016; Guo et al., 2017; Lian et al., 2018). In contrast, the dimension of the embedding matrix grows linearly with the number of categorical features, which can easily be at the scale of millions (Park et al., 2018). As a result, the weight matrix of the embedding layer is often responsible for the major memory consumption of a deep learning based recommendation models. For example, the embedding layer of Facebook recommender system contains billions of parameters. Consequently, the embedding layer occupies more than 99.9% memory of the whole model, which can consume hundreds of gigabytes or even terabytes (Park et al., 2018; Ginart et al., 2021). Without compressing the embedding layers, the excessive memory usage of recommendation models is a major obstacle for serving them on-device, where the memory is limited. Traditional embedding compression methods usually put efforts on compacting the embedding matrix (Markovsky and Usevich, 2012; Wang et al., 2017): Low-rank based methods assume the weight matrix has reduced rank that can be decomposed into several smaller matrices (Markovsky and Usevich, 2012). Hashing based methods reduce the number of embedding vectors in the matrix by mapping similar items into a same bucket (Wang et al., 2017). All these methods follow the framework of the standard embedding technique that learns embeddings with equal length for each token.1 However, recent advances demonstrate that assigning a fixed embedding size to all tokens may be suboptimal due to the following reasons (Joglekar et al., 2020; Zhao et al., 2020a,b; Ginart et al., 2021). In the recommendation domain, usually a few head tokens dominate the data, while the majority of tokens (i.e., long-tail tokens) are rarely observed (Park and Tuzhilin, 2008). Since the token's popularity and the importance of its representation to model performance is correlated (Joglekar et al., 2020; Zhao et al., 2020a; Ginart et al., 2021). Thus, when using a fixed embedding size, it may either lose the information of head tokens or waste parameters on long-tail tokens (Kang et al., 2020; Zhao et al., 2020b). We usually choose a large enough embedding size to ensure model performance, which incurs unnecessary memory usage for storing long-tail token's embedding. To overcome the mentioned drawback of embedding with equal length, several recent work proposes to allocate more capacity (i.e., larger embedding size) to important tokens, and less capacity to unimportant ones (Joglekar et al., 2020; Kang et al., 2020; Zhao et al., 2020a,b; Ginart et al., 2021). These work can be roughly divided into two categories. Some work proposes to explicitly scale token's embedding size with its frequency according to heuristic rules designed by human experts (Kang et al., 2020; Ginart et al., 2021). However, such allocation strategy may be suboptimal since the importance of a token is not purely decided by its popularity. Inspired by neural architecture search (NAS), another line of research formulates the embedding size allocation problem as an architecture selection problem, which selects the embedding size for each token from several predefined options (Joglekar et al., 2020; Zhao et al., 2020a,b). Due to the extremely large search space, the search process incurs a significant computational cost. Although the number of parameters in the embedding layer is significantly reduced, these methods still either have large performance drop or introduce significant extra time cost for searching embedding sizes. In this article, we approach the embedding size allocation problem from a pruning perspective. Our work is motivated by the observation that the majority of token's embeddings can be trained with less capacity without impacting model performance (Joglekar et al., 2020). Therefore, during the search phase, instead of selecting from a set of candidate embedding sizes, we prune the dimensions that have the least impact on model performance in token's embeddings to reduce its capacity. Then, we build a multi-size embedding table for training without sacrificing model performance, where the customized size of each token is obtained by transferring the capacity of its pruned embedding. Moreover, we show that the unimportant parameters in the embedding layer can be identified and pruned at initialization, and this significantly reduces the time cost of searching the customized sizes. Consequently, our framework can reduce the memory occupied by the embedding layer during both the training and inference phases without sacrificing model performance. Our contributions are summarized as follows: We rigorously show that the embedding size allocation problem can be converted to a pruning problem. Based on this reformulation, we propose a pruning-based multi-size embedding (PMB) framework to search the customized embedding size for each token. In our framework, during the search process, the embedding layer is pruned without training it. Thus, the time cost of the search process is significantly reduced. Once pruned, we build the multi-size embedding table for training by transferring the capacity of token's pruned embedding. Our framework can reduce the memory occupied by the embedding layer during both the training and inference phases. We show that our framework can match or improve the performance of several recommendation models using significantly less parameters. e.g., for Autoint+ (Song et al., 2019), we show that PME could significantly improve the Logloss and AUC while using 40 × fewer parameters for click-through rate prediction task on the Criteo dataset. 2. Preliminary and problem statement 2.1. Notations We denote matrices with uppercase bold letters (e.g., V), vectors with lowercase bold letters (e.g., v), and scalars with lowercase alphabets (e.g., v). We use Vi, : to represent the ith row of V, and Vi,j to denote the entry at the ith row and jth column of V. We denote the standard L0 norm as ||·||0. The operation V = concat(V1, V2) represents row-wisely concatenating matrix V1 and V2 into a new matrix V. We use ℕ = {0, 1, 2, 3⋯ } to denote the set of all non-negative natural numbers. We use ⊙ to denote the Hadamard product. 2.2. Preliminary Recommender systems involve a massive amount of categorical feature fields, such as userIDs, itemIDs, and the category of items. Let x = [x1; x2; ⋯ ;xM] be an input instance with M feature fields, where xi is the one-hot vector corresponding to the ith field. Suppose the vocabulary size of the ith field is ni, i.e., there are ni unique tokens (i.e., categorical features) in the ith field. For each token xi, it is mapped into a low-dimensional vector vi∈ℝd by vi = Vixi, where Vi∈ℝni×d is the embedding matrix of the ith field and d is the embedding size. For convenience of notations, let V = concat(V1, ⋯ , VM) be the embedding matrix consisting of all tokens' embeddings. Consider a deep learning based recommender system ϕ parameterized by V and Θ, where Θ denotes all other model's parameters excluding those in V. We denote the prediction corresponding to x as ŷ = ϕ(x|V, Θ). We aimed to minimize the loss L(V, Θ; D) 𝔼 E(x, y)~Dℓ(ϕ(x|V, Θ), y) over a dataset D = {(x, y)}, where ℓ is the loss function such as Logloss. 2.3. Multi-size embedding The multi-size embedding framework allows each token in the vocabulary to have embeddings of different sizes (Joglekar et al., 2020; Ginart et al., 2021). By allocating an appropriate size for each token, the multi-size embedding framework can significantly reduce the total number of parameters in the embedding layer while maintaining the quality of learned representations (Joglekar et al., 2020). Although the multi-size embedding has the mentioned advantages over the standard single-size embedding, applying it requires solving the following problem: Suppose there are n tokens in the vocabulary. If the total number of parameters in the multi-size embedding table is limited to no more than a predefined budget k, how to search for the optimal size di of token i under the budget constraint, such that the loss could be minimized as much as possible with the learned di-dimensional embedding vector v^i? We formally define this embedding size allocation problem in Problem 1. Problem 1 (Embedding size allocation problem). Given a maximum embedding size d and a predefined parameter budget k, let the v^i be a di-dimensional embedding representing token i. For element-wise operations between embeddings to work, embeddings of different sizes are padded to equal length d with zeros following by a projection. Namely, the v^i∈ℝdi will be padded with ei trailing zeros such that di+ei = d, leading to a padded vector v^i′∈ℝd. We define d = [d1, ⋯ , dn]. Let V^∈ℝn×d be the single-size embedding matrix consisting of all projected d-dimensional embeddings, i.e., V^i,:=Piv^i′, where Pi∈ℝd×d is a learnable projection matrix associated with token i. The goal of embedding size allocation problem aimed to solve the following optimization problem: (1) mindL(V^*(d),Θ*(d);D), (2) s.t.  V^*(d),Θ*(d)=argminV^,ΘL(V^(d),Θ(d);D), (3) ∑i=1ndi≤k, (4) ∀i∈{1,⋯,n},di∈ℕ,di≤d. Figure 1 illustrates our multi-size embedding framework. The backbone recommendation models in Figure 1 refer to the rest of the model excluding the embedding layer. Although the projected embeddings have the same number of parameters as the uncompressed ones, we will only retrieve and project the embeddings for tokens in the current mini-batch data. As the mini-batch size restricts the number of retrieved embeddings, the memory usage from these additional parameters is negligible when considering the significant reduction in parameter numbers of the multi-size embedding table. Figure 1 The multi-size embedding framework in our article. For element-wise operations to work (e.g., dot-product in factorization machines), the retrieved embeddings are padded to equal length with zeros following by a field-specific projection. Following the studies by Zhao et al. (2020a) and Ginart et al. (2021), in our article, the projection matrix P in Problem 1 is shared between tokens in a same field to learn field-level structures. We note that such approach also has a nice algebraic explanation: the degree of freedom of the token i's representation is limited by di since (5) Pv^i′=[-p1-⋯-pd-]d×d[v1;⋯;vdi;0;⋯;0︸ei]d×1=∑j=1divjpj. In each field, for the token allocated with larger di, the expressive ability of its embedding is stronger since it is represented using more basis from the row space of P. Thus, the multi-size embedding framework illustrated in Problem 1 can control the capacity of each token's representation by allocating different embedding sizes. Solving Problem 1 poses a significant computational hurdle due to the following two reasons. First, in the recommendation domain, the vocabulary size can easily reach the million level (Covington et al., 2016). Second, since the size of embedding could only be integers, the combinatorial nature of this problem leads to an intractable optimization for a large search space. Finding the optimal embedding sizes for millions of tokens from a discrete search space requires a large amount of computational resources. In the next section, we show that this combinatorial optimization problem can be converted to a pruning problem, which can be approximately solved with significantly less cost. 3. Methodology Figure 2 illustrates the overview of our proposed framework. We first search the customized embedding size for each token in a separate search process before training. The key intuition of our proposed method is the optimal capacity of a token that can be obtained by pruning unimportant dimensions in its embedding. In particular, given a standard single-size embedding layer, we prune the dimensions that have the least impact on model performance in token's embeddings to reduce its capacity. Then, the customized size of each token can be obtained by transferring the capacity of its pruned embedding (Section 3.1). We then derive our proposed pruning-based multi-size embedding framework, which prunes the embedding layer at initialization (Section 3.2). In this way, the time cost of the search process is significantly reduced. Figure 2 Overview of PME framework. In practice, a multi-size table is implemented as multiple two-dimensional embedding matrices, each with different sizes. Since the searched size could be any integer smaller than the maximal size d, we need to initialize at most d two-dimensional matrices, which incurs extra time cost to the retrieval process. To reduce the extra time cost of retrieving from multi-size table, we optimize the retrieval process based on group-wise operations (Section 3.3). 3.1. Size allocation as a pruning problem The success of multi-size embedding framework suggests the embeddings of long-tail tokens can be trained with less capacity without impacting model performance (Joglekar et al., 2020; Ginart et al., 2021). This implies that there exists redundant parameters in the single-size embedding. It is intuitive to start pruning from the parameters that have the least impact on model performance, which is equivalent to reducing the embedding size. For example, as shown in Figure 3, the second value in embedding v1 is pruned out and set as zero, leading to a d1 = d−1 embedding size in effect. The actual size of the pruned embedding equals the number of remaining parameters. Figure 3 An example to illustrate the pruning-based multi-size embedding. After pruning, we build the multi-size embedding table for training, where the size of each token is set to the number of remaining parameters in its pruned embedding. We note that some tokens may be entirely cutoff from the vocabulary (such as v3, in this example), and they are mapped to unlearnable zero vectors. Informally, by setting token i's allocated size di to the number of remaining parameters, the capacity of its pruned embedding will be transferred to v^i in Problem 1. We formalize this statement by showing under mild assumptions, the optimal solution of Problem 1 can be constructed using the pruned embeddings 2. We first give the definition of redundant parameter identification problem. Problem 2 (Redundant parameters identification problem). Given an overparameterized embedding matrix V∈ℝn×d, the redundant parameter identification problem aims to solve the following constrained optimization problem: (6) minV,Θ,CL(V⊙C,Θ;D), (7) s.t.   C∈{0,1}n×d,||C||0≤k, where C is an auxiliary variable representing binary “gates” that denotes whether a parameter in V is present. k is the parameter budget referring to the number of non-zero entries in V, i.e., the amount of gates being “on”. The redundant parameters can be identified by the zeros (the gates being “off”) in C. Proposition 1 (Proof in Appendix 1). If the projection matrix in Problem 1 is shared between tokens in each field, the optimal solution of Problem 1 can be constructed from one solution to Problem 2. The solution d to Problem 1 can be obtained by setting the size of each token to the number of remaining parameters in its pruned embedding. We note that such constructed d satisfies all constraints in Problem 1. First, according to Equation (7), since there are totally at most k remaining parameters in the pruned embedding matrix, the constructed d meets the budget constraint in Equation (3). Second, the constructed d naturally meets the maximal size constraint in Equation (4) since the number of remaining parameters in the pruned embedding are no more than d. As shown in Figure 3, by Proposition 1 and the above analysis, we build the multi-size embedding table for training, where the customized size of each token equals the capacity of its pruned embedding. In the next subsection, we show that Problem 2 can be approximated solved with significant fewer costs. 3.2. Prune embeddings without training them Most of the existing methods in the pruning literature attempt to identify redundant parameters from a pretrained reference network either based on a saliency criterion (Han et al., 2016; Kusupati et al., 2020) or utilizing sparsity enforcing penalties (Carreira-Perpinán and Idelbayev, 2018). Unfortunately, all these pruning methods require many expensive pretrain-prune-retrain cycles and introduce additional hyperparameters. Recent work has explored the possibility of pruning neural networks at initialization (Lee et al., 2019; Wang et al., 2020). Namely, given a desired parameter budget, redundant parameters are pruned once before training, and then the pruned network is trained in the standard way. Equipped with the technique, there is no need for network pretraining and complex pruning schedules. Inspired by single-shot network pruning (SNIP) (Lee et al., 2019), we directly prune unimportant parameters in the embedding according to the connection sensitivity, which can be obtained by utilizing a full-batch of training data. Consequently, the pruning process is disentangled from the above iterative cycle. The key idea of connection sensitivity proposed in SNIP is to preserve the parameters that have the maximum impact on the loss if perturbed. Specifically, the effect of removing parameter Vi, j on the loss can be measured as follows: (8) ΔLi,j(V,Θ;D)=L(1⊙V,Θ;D)-L((1-eij)⊙V,Θ;D), where eij∈ℝn×d is an indicator matrix of element Vi, j (i.e., zeros everywhere except at the ith row and jth column where it is one), and 1∈ℝn×d is an all-ones matrix. Equation (8) measures the influence of parameter Vi, j on the loss in the discrete setting since C is binary. Computing ΔLi, j for each i, j is prohibitively expensive since it requires an individual forward pass over the dataset for each parameter Vi, j. However, by relaxing the binary constraint of C, ΔLi, j can be approximated by the derivative of L with respect to Ci, j, which is named as connection sensitivity. Specifically, the connection sensitivity G(V, Θ; D) in SNIP can be computed as follows: (9) ΔLi,j(V,Θ;D)≈Gi,j(V,Θ;D)=∂L(C⊙V,Θ;D)∂Ci,j|C=1 (10) =∂L(V,Θ;D)∂V⊙V. Parameters that least impact the performance if removed can be identified according to connection sensitivity. We list the full algorithm in Algorithm 1. There is only one hyperparaemter in Algorithm 1, namely, the parameter budget k, which controls the total number of parameters in the multi-size table. Specifically, we first initialize a standard single-size embedding layer, then calculate the connection sensitivity G(V, Θ; D). Once G(V, Θ; D) is obtained, the parameters corresponding to the top-k values of |G(V, Θ; D)| are kept. Finally, the allocated size of each token is set to the number of kept dimensions in its pruned embedding. Algorithm 1 Pruning-base embedding size search. 3.3. Multi-size table lookup optimization Most of the deep learning frameworks do not support embedding table with multiple sizes. In practice, a multi-size table is implemented as multiple two-dimensional matrices, each with different sizes. When retrieving embeddings from a multi-size table, it requires to identify which matrix contains the token's embedding according to its size. The time cost for identifying the matrix containing the token's embedding grows linearly with the number of candidate matrices. In Algorithm 1, the searched size of each token can be arbitrary integer between 0 and d, which means we need to initialize at most d two-dimensional matrices. Thus, the retrieval process will be significantly slowed down when d is large, which contradicts with the goal of being efficient. Similar to the previous studies, (Joglekar et al., 2020; Zhao et al., 2020a,b), we define a candidate size set C={d^1,d^2,⋯,d^T}, where 0≤d^1