The trade-offs of model size in large recommendation models : A 10000 × compressed criteo-tb DLRM model (100 GB parameters to mere 10MB)
Abstract
Embedding tables dominate industrial-scale recommendation model sizes, usingup to terabytes of memory.
A popular and the largest publicly available machinelearning MLPerf benchmark on recommendation data is a Deep Learning Recommendation Model (DLRM) trained on a terabyte of click-through data.
It contains100GB of embedding memory (25+Billion parameters).
DLRMs, due to their sheersize and the associated volume of data, face difficulty in training, deploying forinference, and memory bottlenecks due to large embedding tables.
This paperanalyzes and extensively evaluates a generic parameter sharing setup (PSS) forcompressing DLRM models.
We show theoretical upper bounds on the learnablememory requirements for achieving ( 1 ± ϵ ) approximations to the embedding table.
Our bounds indicate exponentially fewer parameters suffice for good accuracy.
Tothis end, we demonstrate a PSS DLRM reaching 10000 × compression on criteo-tbwithout losing quality.
Such a compression, however, comes with a caveat.
It requires 4.5 × more iterations to reach the same saturation quality.
The paper arguesthat this tradeoff needs more investigations as it might be significantly favorable.
Leveraging the small size of the compressed model, we show a 4.3 × improvementin training latency leading to similar overall training times.
Thus, in the tradeoffbetween system advantage of a small DLRM model vs. slower convergence, weshow that scales are tipped towards having a smaller DLRM model, leading tofaster inference, easier deployment, and similar training times.
Introduction
Recently, recommendation systems have emerged as one of the largest workloads in machine learning[1].
Recommendation systems form the backbone of a good user experience on online platforms suchas e-commerce and web search, where there is a flood of information.
Thus, considerable effort goesinto building recommendation systems.
Deep learning recommendation models give a state-of-the-artperformance.
However, recommendation models suffer from a critical challenge - sparse featureswith millions of categorical values[2, 3] These state-of-the-art [2, 4, 5, 6, 7, 8] methods learn a denserepresentation of the categorical values in a parameter structure called embedding table .
Most parameters in recommendation models come from embedding tables.
For example, in thepopular criteo-tb MLPerf benchmark model, the embedding tables are around 100GB, whereas otherparameters only amount to 10MB.
Industrial-scale recommendation models are one of the largestmodels built.
The size of the embedding table can go as large as hundreds of terabytes.
For example,a research article from Facebook discusses the training of a model of size 50TB over 128 GPUs [3].
The scale of these models leads to some unfavorable effects - slower inference time, slower trainingtime per iteration, and significant engineering challenges in training/deployment.
In light of theseissues, many works have investigated learning of compressed representation of embedding tables using various principles : (1) compositional embeddings [9] (2) exploiting power-law in the observedfrequencies of tokens [10, 11, 12, 13, 14, 15] (3) low-rank decomposition [16], and weight sharingmethods [17, 18].
Parameter-sharing methods will be the focus of our paper.
The general idea in machine learning has primarily shifted to larger models and more data.
Largermodels lead to better capacity, faster convergence, and better generalization.
However, the communityis realizing that the current route to DL success is unsustainable[19].
It stands to reason that we mustproceed cautiously on this path of building larger models.
The large-scale nature of DLRM comesfrom blown-up embedding parameters - a seemingly inadvertent effect of naive usage of embeddingtables.
This paper evaluates the tradeoffs between smaller models achieved by parameter sharingmethods and large models using naive embedding tables.
We find that compressed DLRM modelshave seemingly no downside in training efficiency and quality.
This paper provides a theoretical answer to the question, "How much can embedding tables be compressed without significantly losing quality.
" We present a general parameter sharing setup (PSS) forembedding tables and prove that we only require memory logarithmic (O( 1 / ϵ 2 max ( d 2 ,d. log ( n ))))in n to approximate an embedding table E ∈ R n × d up to a factor of ϵ for subsequent computations.
This result motivates us to evaluate the extents of compression that can be achieved for embeddingtables.
We find that we can obtain 10000 × compression (10MB embedding tables) on the criteo-tbdataset to achieve the same target quality - an order of magnitude higher compression than previouslyreported.
These small models improve inference times and training time per iteration, lower trainingcosts, eliminate engineering challenges and make the models easy to deploy on low-resource devices.
While large compression is possible in theory and as we will show in our results, one important aspectof training compressed models needs to be solved - slow convergence.
A general observation in machine learning is that the number of iterations required to converge to atarget quality decreases with the increase in parameters.
We also observe this effect in the training ofparameter-shared recommendation models.
The higher the compression, the slower the convergence:a concern since training times are one of the significant aspects of the model.
However, a highlycompressed model completely changes the storage of model parameters (e.g., multiple nodes to asingle node, RAM to GPU, etc.).
Hence, we observe a steep decrease in the latency of embeddinglookup/gradient update operations.
An illustration of this tradeoff is presented in figure 1.
Thequestion of how this improvement in latency compensates for increased training iterations is worthinvestigating.
Unfortunately, existing parameter-sharing methods do not demonstrate much trainingtime per iteration improvement due to various reasons such as high algorithmic complexity, poorcache efficiency, and sub-optimal implementation strategies.
This paper talks about various aspectsof implementation.
With our PSS, we demonstrate that, indeed, the advantage of faster training time per iteration almost compensates for the slower convergence in compressed models.
For example a1,000 × compressed model trains in 1 .
1 × time whereas a 10,000 × compressed model trains at 1 . 04 ×time as compared to original model.
This paper provides a theoretical peek into the compression of large recommendation models.
Also,it resolves one of the major roadblocks in training compressed recommendation models by leveragingthe system advantage of compressed models.
We show 10000 × compression ( 10 × higher than bestknown) on DLRM on criteo-tb, making the total model size 20MB.
Such a small DLRM can now beeasily deployed on client devices that are resource constrained.
Theory of parameter-shared embedding tables2.1 General parameter sharing embedding table setup
Definition 1 ( Parameter shared setup ) .
Under parameter shared setup (PSS) for embedding tableof size n × d , we have a set of weights M and recovery function M ∶ R ∣ M ∣ × { 0 ,..,n − 1 } → R d .
Theembedding of a token i ∈ { 0 ,..,n − 1 } is recovered as M( M,i )
M can be represented with any memory layout(eg.
2D/1D array).
We denote the size of M as ∣ M ∣and memory to store M as ∣M∣ .
Note the inherent trade-off between ∣ M ∣ and ∣M∣ .
We can make Mlarge in order to get a very simple M (for example, M = E and M( M,i )=
M [ i ] ) or we can keep Mvery small at the cost of complicated M (for example M =
[0,1] and M combines these bits to get bitrepresentation of E ).
An effective PSS for embedding table E ∈ R n × d would have (∣ M ∣+∣M∣) ≪ nd .
In context of training a PSS in an end-to-end manner, we additionally want (1) M is differentiable(2)
M( M,i ) computation is cheap.
One of the key questions we want to answer is how small a PSS can we create for a ground truthembedding table E ?
i.e. how small can (∣ M ∣+∣M∣) be.
For an arbitrary E , it is impossible to encode
E in sublinear memory accurately.
So, in the next section, we define a ( 1 ± ϵ ) PSS2.2 ( 1 ± ϵ ) parameter sharing setup for table E
While creating an alternative representation of embedding table E , let us see what we would like toachieve.
First, we would like to preserve all the pairwise inner products between two embeddingsof tokens from the embedding table.
Additionally, we also want to minimally affect the subsequentcomputations performed on the embeddings retrieved from the embedding table.
In recommendationmodels, the embeddings retrieved are passed through neural network layers, which will perform innerproducts on embeddings.
With this in mind, we propose the following definition of approximate PSSfor embedding the table.
Definition 2 ( ( 1 ± ϵ ) parameter sharing setup ) .
A PSS ( M, M) is an ( 1 ± ϵ ) PSS for an embeddingtable E if the following holds,
[Equation]
We note that the ( 1 ± ϵ ) PSS for embedding table E also preserves the pairwise distances between theembeddings of different tokens, as is mentioned in the following result.
Theorem 3.
Let ( M, M) is a ( 1 ± ϵ ) PSS for embedding table E, we have ,
[Equation]
Proof.
Presented in the appendix.
We first analyze the norm of recovered embedding from PSS andthen extend the result to pairwise distances.
Naive application of JL-Lemma will give us an existence proof of memory M of size n log ( n ) .
Inthe next section, we show the existence of M with size dlog ( n ) , which is much smaller.
Also, notethat the condition for ( 1 ± ϵ ) PSS is quite different from that of l2-subspace embedding [20].
As we shall see in the proof, this leads to a logarithmic dependence on n (whereas l2-subspace embeddings for A ∈ R n × d has no dependence on n) Also, as we shall see, the considerations in ( 1 ± ϵ ) PSS are the problems
2.3 ( 1 ± ϵ ) PSS via Johnson–Lindenstrauss Transforms
In this section, we answer the question of how small the ∣ M ∣ and ∣M∣ can be.
In order to do this, wegive a construction of a ( 1 ± ϵ ) PSS using the Johnson–Lindenstrauss Transforms(JLT) matrices.
Wewill use the following definition of JLT matrices from [20]
Definition 4.
A randomly generated matrix S ∈ R k × n is a JLT ( ϵ,δ,f ) if with probability at least( 1 − δ ) , for any f -element subset V ⊂ R n the following holds,
[Equation]
The construction of ( 1 ± ϵ ) PSS using JLT is presented in the following theorem,
Theorem 5.
Let the embedding table be E ∈ R n × d .
Consider a matrix S which is
JLT ( ϵ / σ ( E ) ,δ, 9 d + n )) where σ ( E ) denotes the maximum singular value of E. Then M and
M defined by
[Equation]
is a ( 1 ± ϵ ) PSS with probability ( 1 − δ ) where e i ∈ R n is a one-hot encoding of integer
i (i.e e i [ i ] =
1and rest all elements of e i are 0)
Proof.
The complete proof is presented in the appendix.
The intuition is that we want to maintainthe computations of the form e ⊺ i Ex , which can be seen as an inner product between e i ∈ R n and
Ex ∈ R n .
So the inner products between vectors of sets
A = { e 0 ,...,e n − 1 } and column space of Eneed to be preserved.
Once we observe this, we apply JLT to the set of n discrete points and the1 / 2 − net of column space of E which contains less than 9 d points, and get our result.
Note that when using a JLT matrix S as described in theorem 5, the storage cost of memory M is∣ M ∣ =
∣ SE ∣ =
kd .
For the mapping M , we only need to store the matrix S .
Thus ∣M∣ =
∣ S ∣ .
The costof storing S will depend on the exact choice of JLT.
Since ( Se i ) is a O( 1 ) operation as e i is a one-hotencoded vector, the cost of M( M,i ) is just that of vector-matrix multiplication between ( Se i ) and
M .
Again, the exact cost depends on the choice of S. Using the random matrix formed with i.i.dnormal variables, the most basic form of JLT, we can give the following bound on the memory oflearnable parameters M
Corollary 6.
Let E be the embedding table under consideration.
Also, let S ∈ R k × n matrix whereeach element of S is i.i.d N( 0 , 1 k ) with k = Ω ( σ 2 ( E )/
ϵ 2 ( log ( 1 / δ ) + max ( d,log ( n ))) , then with aprobability ( 1 − δ ) , ( M = SE, M( M,i ) =
( Se i ) ⊺ ( M )) is a ( 1 ± ϵ ) PSS for E.
Thus we obtain an upper bound on M. ∣ M ∣
=
O
( 1 / ϵ 2 max ( d 2 ,d.log ( n )))
when using the standardnormal matrix as JLT for embedding tables with bounded singular values.
In this case, the mappingstorage cost is ∣M∣ =
kn as the matrix S is dense.
The cost of applying M will be kd and hence,
O( 1 / ϵ 2 max ( d 2 ,d.log ( n ))) .
In the next section, we will see how to reduce the cost of storing Sin practice.
This logarithmic dependence on n in the case of embedding tables with very large nexplains why we can achieve a very high compression of 10000 × for the criteo-tb dataset.
2.4 ( 1 ± ϵ ) PSS in practice( 1 ± ϵ ) PSS from a learned embedding table E
The above discussion directly gives us an algorithmto compute a ( 1 ± ϵ ) of a trained embedding table E .
We just need to compute M = SE and whileretrieving an embedding compute M( M,i ) =
( Se i ) ⊺ M .
Training end-to-end ( 1 ± ϵ ) PSS for embedding table E We can also directly train an ( 1 ± ϵ )
PSSin an end-to-end manner.
The idea is to directly train the compressed M .
Thus, we have a matrix oflearnable weights
M ∈ R k × d .
Let us now look at how the forward and backward pass of embeddingretrieval mapping M looks like
[Equation]
The forward function takes an integer i and returns a R 1 × d array which is embedding of i .
[Equation]
The backward pass takes in all the arguments of the forward pass along with the ∆ ∈ R 1 × d which isgradients of loss with respect to the output of the forward pass.
We can back-propagate further to Wusing the above formulation.
The result of the backward pass is a k × d matrix.
If we use a sparse S such as with sparse JLT, we can achieve sparse gradients of W. That is, only afew entries of the gradient for W are non-zero.
As we will see in section 3, whether to propagatesparse or dense gradients is a vital implementation choice.
We were able to achieve better trainingper iteration by exploiting this choice.
Cost considerations for PSS using JLT.
In using PSS in end-to-end training, we would be learningthe memory M while keeping the mapping M constant.
Thus, unlike in sketching for linear algebra[20], we care about the costs of (1) storage of learnable parameters ∣ M ∣ , (2) Cost of storing mapping∣M∣ , and (3) cost of computing M( M,i ) .
It is clear that using standard normal JLT is not feasibledue to storage of matrix S will be more expensive than storing E itself!
There are a lot of sparse
JLT ( ϵ,δ,f ) matrices S which will reduce the storage cost of S [21, 22, 23].
[23] showed that thematrix needs to have a minimum column sparsity of Ω ( ϵ − 1 log ( f / δ ) log ( 1 / ϵ )) and hence the costof storing S while using JLT matrices is lower bounded by Ω ( nϵ − 1 log (( n + 9 d )/ δ ) log ( 1 / ϵ )) whichcan still be considerably high.
In the next section, we see relaxations of conditions on S for practicalpurposes.
2.5 Practical PSS and existing SOTA methods
This section will briefly review practical relaxations of PSS defined above.
Sparse JL with independence relaxation: The issue with storage costs of S can be alleviated inpractice by relaxing the complete independence condition on entries of S. For example, the JLTmatrix (which requires the same bound on k as given in theorem 5) suggested by [21], which selectseach entry to be 0 with probability 2/3, ± 1 with probability 1/6 independently can be generatedon-the-fly using universal hashing.
[24] analyze why simpler hash functions can work well with datahaving enough entropy.
The benefit of using universal hashing is that mapping M can now be storedin O ( 1 ) memory.
Thus using the [21] sparse matrix, we can obtain ( 1 ± ϵ ) PSS using total setupmemory of ∣ M ∣ + ∣M∣ =
O( 1 / ϵ 2 max ( d 2 ,d.log ( n ))) and cost of applying mapping M to be O( d )(hash computation and lookups) under valid independence relaxation.
Existing parameter-sharing based embedding compression methods can also be seen as PSS withvaried distributions over S.
We state a few of them below.
• Hashing Trick
In this method, entire embeddings for a token i is drawn from a randomly hashed
location.
In terms of PSS, we can define the mapping function M h over a 2D matrix M h ∈ R k × dwhere k < n .
[Equation]
where h is a hash function.
• QR decomposition [9].
In this method, embedding for a token i is drawn in chunks from separatememory vectors.
In terms of PSS, we can define the mapping function M q over, say l, pieces of2D memory M 1 ,M 2 ,..
M l ∈ R k × d / l as follows
[Equation]
In words, the j th chunk of i th embedding is recovered using the chunk at the location h j ( i ) in
memory M r
[Table]
• HashedNet [17] and ROBE-Z [18]
: In HashedNets, authors proposed mapping model weightsrandomly into a parameter array.
ROBE-Z extended this idea by hashing chunks of embedding
vector instead of individual elements.
In terms of PSS, we can define the mapping function M r
over 1D memory M r as
[Equation]
In words, the j
th chunk
of i th embedding
is recovered using the chunk at the location h ( i,j ) from
M r .
Here, h ∶ N 2 → 0 ,.., ∣ M r ∣} is a
hash
function
If we set Z = 1 ,
then
we get the
mapping
function for HashedNet.
We can summarize the state-of-the-art embedding compression in the table 1.
Our PSS theorysuggests that we should be able approximate the embedding table in memory of the order of O (log(n).
Thus, it is important to question if 1000 × compression is the best compression we can achieve on adataset like CriteoTB where we have large ( ∼ 400GB sized) embedding tables.
As we shall see, oneof the significant roadblocks in aiming to achieve higher compression is training time.
The higher thecompression, the more iterations are needed to train the model to a certain quality.
In the next section,we discuss how to overcome this bottle-neck preventing training of highly compressed PSS.
Road to 10000 × compression with PSS
Table 4 shows the quality of embeddings for the criteo-kaggle dataset over varying compression ratesacross different popular recommendation models.
We can see that the state-of-the-art compression ofthe criteo-kaggle dataset is 1000 × .
It is a consequence of PSS theory that larger embedding tables(prevalent in industrial scale recommendation models) such as those for criteo-tb should obtain morecompression.
Our goal is to study higher compression for large embedding tables and improve overexisting SOTA ( 1000 x).
A significant roadblock in building highly compressed embedding tables is the slow convergence.
Table 4 shows that across different deep learning recommendation models, the number of iterationsrequired to converge increases with higher compression.
Specifically, for 1000 × compression, it takesup to 4 × iterations, and for models with 10000 × compression, most models do not converge in 15epochs.
If a model does not converge, it is hard to judge if the convergence is too slow or its capacityis consumed.
In order to be able to train for larger iterations, we should be able to train the highlycompressed models faster.
Indeed, we expect some system benefits with highly compressed models,as shown in figure 1.
We find that existing methods cannot train highly compressed embedding tablesbecause of the implementation choices.
We propose the following implementation choices.
Hashing chunks is better than hashing elements: As suggested in [18], we find that hashing chunksinstead of individual elements is indeed helpful in reducing the latency of PSS systems.
As shown intable 3(d,e), both forward and backward passes are affected a lot by chunk sizes.
It is also interestingto note that the effect of chunk sizes diminishes after size 16.
Hence, we can choose a chunk size of32 for hashing.
Dense gradients are suitable for highly compressed PSS Generally, in embedding tables, implementations use sparse gradients.
Dense gradients are particularly wasteful in embedding tables,where in each iteration, only a few weights are involved in the computation.
However, the samedoes not apply to highly compressed PSS.
In the case of PSS, the algorithm for sparse gradientcomputation is more involved than that for dense gradients.
The algorithms for sparse and densegradient computation are specified in algorithm 1.
Also, the scatter operation is much faster when thememory size of embeddings is smaller.
Hence, dense gradients work well with higher compressions.
From table 3(b,c), we can see that at higher compression rates, the backward pass is up to 8 × fasterwith dense gradients as compared to sparse gradients.
[Figure]
[Table]
Require:
G
O ∈ R B × d , I ∈
R B × d , sparse ∈
{ True , False } , M: compressed memory.
▷ G O is the gradient w.r.t output, and I is the mapping matrix showing the locations from which the output was accessed
Ensure: G i is gradient w.r.t input.
if sparse then unique , invIdx ← findUnique (I . flatten () , returnInverse = True )
g ← Tensor
( unique . shape , ))
g. scatter _ add ( invIdx , G o .
flatten ()
G i = sparseTensor ( location = unique , values = g , size = M . size ) else if dense then
G i ← Tensor ( M . size )
G i .
scatter _ add ( 0 , I .
flatten () , G o .
flatten ()) end if
Forward/Backward kernel optimizations Latency for PSS, and CUDA kernels in general, is verysensitive to the usage of shared memory, occupancy, and communication between CPU and GPUs.
We also optimize our PSS code to minimize the data movement costs, implement custom kernels tofuse operations to improve shared memory usage and optimize CUDA grid sizes to obtain the bestperformance.
With the implementation choices mentioned above, we are ready to train a 10000 × compressed PSSfor criteo-tb.
We implement our PSS using ROBE-style hashing.
Experimental results on Criteo datasets
We perform experiments on criteo-kaggle and criteo-tb datasets in order to confirm the followinghypothesis,
Theory in section 2 dictates that
embedding table ( E ∈ R n × d ) can be represented in memorylogarithmic in n .
High compression should be possible in these datasets.
Specifically, highercompression should be possible in larger embedding tables.2.
The system advantage of smaller PSS should compensate for the convergence advantage of theoriginal model.
Datasets: Criteo-kaggle and criteo-tb datasets have 13 integer and 26 categorical features.
criteokaggle data was collected over seven days, whereas the criteo-tb dataset was collected over 23 days.criteo-tb is one of the largest recommendation datasets in the public domains with around 800 milliontoken values making the embedding tables of size around 400GB.
(with d = 128).
criteo-kaggle issmaller and has a 2GB sized embedding table.
Note that industrial-scale models are much larger than
[Figure]
deviation
of AUC results
is within
[Table]
a) Epochs required to reach 0.8025 AUC for DLRM model on criteo-tb dataset.
Relative time computes the ratio of A*B to original model
[Table]
) 10 4 × compression also reaches the target
AUC.
For 10 5 × convergence is too slow/mode capacity is reached. (max 15 epochs
[Table]
the DLRM model we talk about here.
For example, Facebook recently published a model sized 50TB[3].
One can extrapolate the benefits of PSS to industrial-scale models based on this case study.
Models: Facebook MLPerf DLRM[2] model, available under Apache-2.0 license, for the criteo-tbdataset achieves the target AUC (0.8025) with the embedding memory of around 100GB.
This modeluses a maximum cap of 40M indices per embedding table, leading to a total of 204M embeddings.
This model cannot be trained on a single GPU (like V100) and is trained using a multiple GPUs (4 or8).
For the criteo-tb dataset, we use the DLRM MLPerf model.
For criteo-kaggle dataset, we use anarray of state-of-the-art models DLRM[2], DCN[4], AUTOINT[5], DEEPFM [6], XDEEPFM [8]and FIBINET.
We use our PSS implementation described in 3 as a compressed model.
Quality of model vs. Excess parameters: Tables 4(b) and 5(b) show the results for the two datasetsacross different values of compression.
In table 4(b), we can see that across different models, thequality of the model is maintained until 1000 × compression.
As criteo-tb embedding tables aremuch larger than the criteo-kaggle dataset, according to the section 2, we should see larger valuesof compression.
Indeed, we obtain a 10000 × compression without loss of quality of the model forcriteo-tb.
This level of compression is unprecedented in embedding compression literature.
Theseexperiments validate our first hypothesis and provide a new state-of-the-art embedding compression.
In the rest of the section, we evaluate how the system advantage of PSS compares against the originalembeddings in various aspects.
Inference time In figure 6a, we compare the inference time of test data (89M samples) in the criteo-tbdataset with a batch size of 16384 with PSS on a single Quadro RTX-8000 GPU and the originalmodel on 8 Quadro
RTX-8000 GPUs.
The total time for inference for the original model is around203 seconds, whereas PSS is around 97 seconds.
In the original model trained on 8GPUs, embeddinglookup is model-parallel, whereas the rest of the computation ( bot-mlp, top-mlp, and interactions) isdata-parallel.
There is a steep improvement of over 3 × in embedding time lookup as we move from adistributed GPU setup for embedding tables to a PSS on a single GPU.
The extra computation time in
PSS is much smaller than all-to-all communication costs in the original embedding table lookup.
Weinclude the time required for data distribution (initial) in bot-mlp and data-gather (final) in top-mlptimings.
This communication cost is high, and we can see that we are better off performing the entirecomputation on a single GPU.
Model training time per iteration vs excess parameters
Figure 6b records the time-taken for 1000iterations of training.
While using embedding tables, we can choose to back-propagate sparse ordense embedding gradients.
The general idea is to use sparse gradients when very few gradientsare non-zeros.
In the original model, we can only use sparse gradients as using dense gradients isprohibitive w.r.t to computation and communication.
In PSS, however, we compare both modes forgradient back-propagation.
We see that dense gradients perform exceptionally well at higher rates ofcompression.
The performance of sparse gradients is constant across different compression rates asthe workload is similar.
Generally, it seems good to use dense gradients for PSS when the effective memory size (i.e., the final memory of compressed embedding tables ) is small.
It is noteworthythat the training time per iteration reduces significantly with higher compression.
For example, with1000x compression, the training time is 1 .
7 × lesser than the time taken by the original model,
whereaswith 10000x compression, the time per iteration is 4.37 × lesser.
Model convergence and overall time vs. excess parameters We observe a consistent trend inconvergence: as the number of excess parameters in the models reduces, the convergence becomesslower.
This can be seen across models and datasets as shown in tables 4(a) and 5(a).
As anexample, table 5 shows that with 10000 × compression, we require 4.55 epochs as compared to theoriginal model’s one epoch.
This might seem unfavorable at first.
However, as seen in the previoussection, there is a significant gain in training time per iteration.
As seen in table 5, this gain largelycompensates for the disadvantage in terms of convergence, making the overall training times similar.
For example, while we need 4.55 epochs with 10000 × compressed PSS, the training time per iterationgoes down by a factor of 4.37.
Hence, the overall training time is only 1.04 × original time.
Engineering challenges and costs vs. excess parameters: The excess parameters in a recommendation model primarily appear due to the construction of embedding tables.
The industrial-scaleembedding tables can go as large as hundreds of terabytes.
With the increase in the model’s size,the model needs to be distributed across different nodes and GPUs.
The complexities of efficientlyrunning a distributed model training include many considerations such as non-uniform memoryallocation and communication costs.
The increasing literature detailing the engineering solutions totraining such models is evidence of the fact that training such models require engineering ingenuityto solve the challenges involved [3, 25, 26, 27, 28, 29].
For example, in [3], the authors detail theirsolution to train a model of size 50TB on a new distributed system.
Similarly, in [25], authors talkabout their optimizations of recommendation models on CPU clusters.
In this paper, we argue thatthe large-scale nature of embedding-based recommendation models appears due to embedding tableswhere a complete n × d table with blown up n does not add value to model capacity while addingsignificant engineering challenges.
Distributed systems also imply significant energy costs.
PSS canavoid these downsides by fitting entire learnable parameters on a single GPU/node.  Limitations and future work
In this paper, we show that we can compress the embedding tables for CTR prediction models asmuch as 10000 × .
The theory presented in the paper shows promise of generalizing these resultsto other models such as retrieval or ranking problems.
However, more experiments are needed toevaluate PSS approaches in these problems.
We evaluate CTR models in this paper and extensions ofthis work to other models is left for future work.
Conclusion
Today’s prevalent idea in deep learning(DL) is overparameterized models with stochastic iterativelearning.
While this paradigm has led to success in DL, the sustainability of this route to successis a grim question [19].
We evaluate the value of embedding-based excess parameters in DLRMmodels.
We conclude that compressed DLRM models are better for fast inference and are easy totrain and deploy as expected.
More importantly, they can also be trained in the same overall timedespite having a slower convergence rate.
The critical observation is that compressed models enjoy a system advantage they can exploit, which reduces time per iteration.
The paper also highlights thatthe correct way to compare the quality of models of highly different sizes is to run the models forequal time rather than equal iterations.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] It is clear in theory
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] see section
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes] see supplementary
If you ran experiments...
(a)
Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No] Our code is aset of patches on exisitng codes and hence is difficult to package.
we will publish linkswhen accepted(
b) Did
you specify all the training details (e
.g. data splits hyperparameters, how
they
were chosen)?
[Yes] yes in supplementary
(c) Did you report error bars (e
g., with respect to the random seed after running experi
ments multiple times)?
[Yes] See
table captions
(d) Did you include the total amount of compute and the type of resources used (e
g., type of GPUs, internal cluster, or cloud provider)?
[Yes]
In experimental section
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes] In models section in experiments
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
A AppendixB Proof of Theorem 5
Let the embedding table be E ∈ R n × d .
Consider a matrix S which is JLT ( ϵ / σ ( E ) ,δ, 9 d + n )) whereσ ( E ) denotes the maximum singular value of E. Then M and M defined by
M = SE
M( M,i )=(
Se i ) ⊺ ( M ) is a ( 1 ± ϵ ) PSS
where e i ∈ R n is a one-hot encoding of integer i (i.e e i [ i ] = 1 and rest all elementsof e i are 0)
Proof.
Johnson Lindenstrauss transforms are a class of transformations which preserve pair wiseinner products ( equivalently norms ) of a set of vectors in a space Definition of JLT transform from[20] is A random matrix S ∈ R k × n forms a JLT ( ϵ,δ,f ) if with probability ( 1 − δ ) for any f-elementsubset V ⊂ R n , for all v 1 ,v 2 it holds that
[Equation]
Choice of S matrix and inner product preservation
Consider the following sets of points.1.
{ e i } n
[Equation]
2. 1 / 2 − net over the set { u ∣ u ∈ Column-space ( E ) ∣∣
u ∣∣∣ 2 = 1 }
We use the lemma 5 from [20] to have the number of points in set 2 bounded by 9 d .
Let S be a matrixthat is JLT( δ,ϵ, ( n + 9 d ) ).
For these ( n + 9 d ) points with probability ( 1 − δ ) , the matrix S preservesinner products as per definition of JLT given above.
Let us consider 4 cases w.r.t to the followingcondition
[Equation]
v 1 ,v 2 ∈ Column-space ( E ) .
We use the same argument as given on page 12 [20] andconclude that condition holds for these v 1 ,v 22.
v 1 ,v 2 ∈ { e i } n − 1 i = 0 ∪ 1 / 2 − net .
condition holds due to JLT3.
v 1 ∈ { e i } n − 1 i = 0 ,v 2 ∈ column-space ( E ) .
Let ∣∣ v 2 ∣∣ = 1 .
We will prove the condition for thiscase and for all other non-unit-norm cases will happen due to scaling.
Using argumnet frompage 12 [20], we can represent v 2 as a sum of vectors in 1 / 2 − net .
[Equation]
where v
i
[Equation]
[Equation]
net
[Equation]
4) such that
[Equation]
[Equation]
Thus, using the matrix S, for all v 1 ,v 2 ∈ { e i } ∪ column-space ( E ) it holds that ,
[Equation]
proving the statement for ( 1 ± ϵ ) PSS Let e i s be one hot encoded vectors such that we have that
[Equation]
[Equation]
Let M = SE and M( M,i ) =
( SE )
⊺
Se i
[Equation]
[Equation]
Note that both e i and Ex are vectors that belong to our f point set
V. Thus
[Equation]
Note that ∣∣ e i ∣∣ =
1 Let σ be the max singular value of E. We can push σ ( E ) into ϵ and select atransform JLT ( ϵ / σ ( E ) ,δ, ( n + 9 d )) to obtain a ϵ − l 2 parameter shared setup.
C Proof of theorem
Let ( M, M) is a ( 1 ± ϵ ) PSS for embedding table E, we have ,
[Equation]
i,
[Equation]
2 ϵ
Proof.
Using definition of PSS
[Equation]
and ⟨M( M,i ) , M( M,i )⟩ ≥ ⟨M( M,i ) ,E
[ i ]⟩ − ϵ ∣∣M( M,i )∣∣⟨M( M,i ) , M( M,i )⟩ ≤ ⟨ E [ i ] ,E
[ i ]⟩ + ϵ (∣∣ E
[ i ]∣∣ 2 + ∣∣M( M,i )∣∣ 2 ) and ⟨M( M,i ) , M( M,i )⟩ ≥ ⟨ E [ i ] ,E
[ i ]⟩ − ϵ (∣∣ E
[ i ]∣∣ 2 + ∣∣M( M,i )∣∣ 2 ) ⟨M( M,i ) , M( M,i )⟩ ≤ ⟨ E
[ i ] ,E
[ i ]⟩ + ϵ (∣∣ E
[ i ]∣∣ 2 + ∣∣M( M,i )∣∣ 2 ) and ⟨M( M,i ) , M( M,i )⟩ ≥ ⟨ E [ i ] ,E
[ i ]⟩ − ϵ (∣∣ E
[ i ]∣∣ 2 + ∣∣M( M,i )∣∣ 2 )
Let ∣∣M( M,i )∣∣
= m and ∣∣ E
[ i ]∣∣ = em 2 − ϵm ≤ e 2 + ϵe
and m 2 + ϵm ≥ e 2 − ϵe
{+Adding ( 1 / 4 ) ϵ 2 on both sidesm 2 − ϵm + 1 / 4 ϵ 2 ≤ e 
+}
+
ϵe
{++ ϵ 2 and m 
+}
+ ϵm + 1 / 4 ϵ 2 ≥ e 2 − ϵe + 1 / 4 ϵ
{+( m − 1 / 2 ϵ ) 2 ≤ ( e + 1 / 2 ϵ ) 
+}
and ( m + 1 / 2 ϵ ) 2 ≥ ( e + 1 / 2 ϵ )
[Equation]
and
[Equation]
Thus,
[Equation]
Thus,
[Equation]
Now we can look at the pairwise distances
[Equation]
Thus,
[Equation]
D.1 PSS rigorous evaluation
• Page 1 : memory vs compression.
• page 2: memory vs chunk size
• page 3: compression vs chunk size
Some obervations
• (sparse/dense) original embeddings should be run with sparse gradients.
The dense embed-
dings are too time consuming as they require updating large amounts of memory (vacuously)
in each iteration.
Note that a lot of optimizers in deep learning libraries like pytorch do notyet have support for sparse gradient updates.
• When the entire embedding table can fit on the GPU, the forward and backward times donot change much.
• While we can go upto 64M tokens ( for m = 128 ) on single gpu using original embeddings
and sparse gradient propogation, the simulated embedding tables can be much larger withPSS and compression.
• Best forward times we can achieve with PSS ( around 0.27ms/iteration ) is almost 2 × the
forward times for original embedding lookup.
(0.12 ms / iteration).
This is expected since,
we have to perform additional hash computations in PSS.
Hence, if both original embeddings
and compressed embeddings are at the same distance from the computational resource, then
PSS will have a disadvantage.
• (compression, PSS)
Higher the compression, better are timings for both backward pass
(using dense propagation) and forward pass.
• A nice trade-off can be seen between using sparse or dense gradients with PSS.
If we
use dense gradients, the time in backward propogation is affected by the overall size of
embedding parameters.
Thus, for higher compression and smaller n , we have smaller times
where as if n becomes larger and compression is smaller, the time taken increases.
The cost
of dense gradient propogation becomes quite high at sufficiently large n and sufficiently
small compression.
On the other hand, the cost of sparse gradient propogation is uniform
across different values of n and compression.
This is because the
algorithmic complexityand memory
accessed is similar
This
gives us
a
guideline as to when to use
sparse /
dense
gradients with PSS
Note
that
in
our
final
results, we see a good overall improvement by
leveraging the dense gradient propogation at high compression
• Higher the chunk size better is the time in forward pass.
backward pass is largely unaffected
by chunk size largely due to the way it is implemented.
Original time(ms) n call sparse 4M 8M 16M 32M 64M 128M backward
FALSE 25.
TRUE 0.
forward
FALSE 0.
TRUE 0.
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
dimension =
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]