F REQUENCY - AWARE SGD FOR E FFICIENT E MBEDDING L EARNING WITH P ROVABLE B ENEFITS
Anonymous authors
Paper under double-blind review
A BSTRACT
Embedding learning has found widespread applications in recommendation systems and natural language modeling, among other domains.
To learn quality embeddings efﬁciently, adaptive learning rate algorithms have demonstrated superior empirical performance over SGD, largely accredited to their token-dependent learning rate.
However, the underlying mechanism for the efﬁciency of tokendependent learning rate remains underexplored.
We show that incorporating frequency information of tokens in the embedding learning problems leads to provably efﬁcient algorithms, and demonstrate that common adaptive algorithms implicitly exploit the frequency information to a large extent.
Speciﬁcally, we propose (Counter-based) Frequency-aware Stochastic Gradient Descent, which applies a frequency-dependent learning rate for each token, and exhibits provable speed-up compared to SGD when the token distribution is imbalanced.
Empirically, we show the proposed algorithms are able to improve or match adaptive algorithms on benchmark recommendation tasks and a large-scale industrial recommendation system, closing the performance gap between SGD and adaptive algorithms, while using signiﬁcantly lower memory.
Our results are the ﬁrst to show token-dependent learning rate provably improves convergence for non-convex embedding learning problems.
I NTRODUCTION
Embedding learning describes a problem of learning dense real-valued vector representation for categorical data, often referred to as token (Pennington et al., 2014; Mikolov et al., 2013a;b).
Good quality embeddings can capture rich semantic information of tokens, and thus serve as the cornerstone for downstream applications (Santos et al., 2020).
Due to their signiﬁcant impact on model performance and large memory footprint ( 21 . 8% of total parameters for BERT (Devlin et al., 2018), 95% for industrial recommenders in Section 4), how to learn quality embedding vectors efﬁciently forms an important problem in applications, including recommendation systems and natural language processing.
Empirically, adaptive algorithms (Duchi et al., 2011; Kingma & Ba, 2014; Reddi et al., 2019) have witnessed signiﬁcant successes, yielding state of the art performance in both industrial-scale recommendation systems and natural language model (Guo et al., 2017;
Zhou et al., 2018b; Devlin et al., 2018; Liu et al., 2019).
Stochastic gradient descent (SGD), on the other hand, has struggled to keep up, often yielding much slower convergence and low quality models (Liu et al., 2020; Zhang et al., 2019)
(see also Figure 2).
The sharp contrast on the efﬁciency of adaptive algorithms and SGD is particularly distinctive, as SGD is the typical choice of optimization algorithms in the other domains of machine learning, such as vision/image related tasks (He et al., 2016; Goyal et al., 2017).
The common belief behind the empirical edge of adaptive learning rate algorithms over SGD is that the former ones exploit sparsity of high dimensional feature.
Speciﬁcally, a feature in a typical embedding learning problem comes in the form of one/multi-hot encoding of tokens (e.g. workpiece in NLP and user/item in recommendation systems), which leads to a sparse stochastic gradient that has only non-zero values for tokens within the mini-batch.
In addition, token distributions of real world data are often highly imbalanced and satisfy the power-law property (Piantadosi, 2014; Celma, 2010; Clauset et al., 2009), and infrequent tokens are widely believed to be more informative to model learning.
Thus adaptive algorithms can pick up information from the infrequent tokens more efﬁciently, as they can schedule a higher learning rate for the infrequent tokens (Duchi et al., 2011).
Despite the appealing intuition, there is a signiﬁcant theory-practice gap on the empirical superiority of adaptive learning rate algorithms over SGD, and no developed theories can explicitly justify the previous intuition.
Better dimensional dependence of adaptive algorithms has only been shown in the convex setting (Duchi et al., 2011), which hardly generalizes to even the simplest practical models in embedding learning problems (e.g., Factorization Machine, Rendle (2010)), whose loss landscape is non-convex.
For non-convex settings, most theoretical efforts have been devoted to analyzing adaptive learning rate algorithms for general non-convex objectives, which yield subpar convergence rate compared to standard SGD (Ward et al., 2018; D´efossez et al., 2020; Chen et al., 2018;
Zhou et al., 2018a).
In fact, the standard SGD has been recently shown to be minimax optimal for non-convex problems (Drori & Shamir, 2020; Arjevani et al., 2019), and thus not improvable in general.
Moreover, since adaptive algorithms are only implicitly exploiting frequency information, and if the intuition indeed holds true, one might naturally wonder whether we can instead develop an adaptive learning rate schedule that explicitly depends on frequency information.
Motivated by our previous discussions, we raise and aim to address the following questions
Questions
Can we design a frequency-dependent adaptive learning rate schedule?
Can we show provable beneﬁts over SGD?
Our contributions.
We answer the previous question by showing that token frequency information can be leveraged to design provably efﬁcient algorithms for embedding learning.
Speciﬁcally,
• We propose F requency- a ware S tochastic G radient D escent (FA-SGD), a simple modiﬁcation to standard SGD, which applies a token-dependent learning rate that inversely proportional to the frequency of the token.
We also propose a variant, named C ounter-based F requency-aware S tochastic G radient D escent (CF-SGD), which is able to estimate frequency in an online fashion, much similar to Adagrad (Duchi et al., 2011) and Adam (Kingma & Ba, 2014).
• Theoretically, we show that both FA-SGD and CF-SGD outperform standard SGD for embedding learning problems.
Speciﬁcally, they are able to signiﬁcantly improve convergence for learning infrequent tokens, while maintaining convergence speed for frequent tokens.
To the best of our knowledge, our proposed algorithms are the ﬁrst to show provable speed-up over standard SGD for non-convex embedding learning problems.
This is in sharp contrast with other popular adaptive learning rate algorithms, whose empirical performance can not be explained by existing theories.
• Empirically, we conduct extensive experiments on benchmark datasets and a large-scale industrial recommendation system.
We show that FA/CF-SGD is able to signiﬁcantly improve over SGD, and improves/matches popular adaptive learning rate algorithms.
We also observe the secondorder moment maintained by Adagrad and Adam highly correlates with the frequency information, demonstrating intimate connections between adaptive algorithms and the proposed FA/CF-SGD.
1. R ELATED L ITERATURE
Adaptive algorithms for non-convex problems.
There has been a fruitful line of research on analyzing the convergence of adaptive learning rate algorithms in non-convex setting.
These results aim to match the convergence rate of standard SGD given by O (1 / √ T ) (Ghadimi & Lan, 2013), however often with additional factor of log T (Ward et al., 2018; D´efossez et al., 2020; Chen et al., 2018; Reddi et al., 2018), or with worse dimension dependence (Zhou et al., 2018a) for smooth problem (assumed by almost all prior works).
Moreover, all existing works aim to analyze the convergence for general non-convex problems, ignoring unique data features in embedding learning problems, where adaptive algorithms are most successful.
We explicitly take account into the sparsity of stochastic gradient, and token distribution imbalancedness into the design and analysis of our proposed algorithms, which are the keys to better convergence properties.
Adaptive algorithms and SGD.
To the best of our knowledge, the study on understanding why adaptive learning rate algorithms outperform SGD is very limited.
Zhang et al.
(2019) argue that
BERT pretraining (Devlin et al., 2018) has heavy-tailed noise, implying unbounded variance and possible non-convergence of SGD
Normalized gradient clipping method is proposed therein and converges for a family of heavy-tailed noise distributions.
Our results focus on a different direction by showing that imbalanced token distribution is an important factor that can be leveraged to design more efﬁcient algorithms for embedding learning problems.
Our result also does not rely on the noise to be heavy-tailed for the convergence beneﬁts of the proposed FA/CF-SGD to take effect
[Equation] [Equation]
Notations: For a vector/matrix, we use (cid:107)·(cid:107) to denotes its (cid:96) 2 -norm/Frobenius norm.
We use (cid:107)·(cid:107) 2 to denote the spectral norm of a matrix.
P ROBLEM S ETUP
We consider an embedding learning problem which aims to learn user and item embeddings through their interactions.
We denote U as the set of users, and V as the set of items, and let X
=
U ∪ V denote the union, referred to as tokens throughout the rest of the paper.
We assume | X | = N , i.e., the total number of user and item is N .
For the ease of presentation, we always use letter i to index user set U , letter j to index item set V , and letter k to index the union set X .
The embedding learning problem can be abstracted into the following stochastic optimization problem: min
[Equation]
Here ( i, j ) denotes the user-item pair sampled from the unknown interaction distribution D , θ i , θ j ∈ R d (the i, j -th row of Θ ) denotes their embedding vectors respectively, and the loss (cid:96) ( θ i , θ j ; y ij ) denotes the prediction loss for their interaction y ij ∈ {− 1 , +1 } (e.g., logistic loss).
We further let
[Equation]
denote the marginal distribution over U and V .
Remark 2.1 .
Our analysis also allows treatment of additional network structure (with parameters denoted by W ) that takes nonlinear transformation of embedding vectors, e.g., f (Θ , W )
=
E ( i,j ) ∼D (cid:96) ( θ i , θ j , W ; y ij ) .
We omit their explicit treatment for presentation simplicity.
In addition, although we mainly discuss in the context of recommendation, our analysis and results only relies on sparsity of stochastic gradient and the imbalancedness of token distributions, which allow one to extend our results to other embedding learning problems (e.g., language model pretraining).
The full algorithmic descriptions of our proposed Frequency-aware Stochastic Gradient Descent (FA-SGD) algorithm are presented in Algorithm 1.
Note that randomly outputting a historical iterate is commonly adopted in literature for showing convergence of stochastic gradient descent type algorithms for non-convex problems (Ghadimi & Lan, 2013).
In practice, we can simply use the last iterate Θ T as the output solution.
In addition, Section 3.3 presents CF-SGD (Algorithm 2), which does not need the token distribution as the input and can estimate it in an online fashion.
[Equation]
At iteration t , FA-SGD samples ( i t , j t ) ∼ D , and obtain the sparse stochastic gradient g t deﬁned in (3).
Note that only the i t -th and j t -th row of g t are non-zero.
One can readily verify that E ( i t ,j t )
∼D [ g t ] = ∇ Θ f (Θ t ) .
Going forward, we will denote ∇ f tk as the k -th row of gradient ∇ f (Θ t ) , and g tk as the k -th row of stochastic gradient g t .
Note that we have
[Equation] [Equation]
We further denote δ tk = 1 p k f tk − g tk for all k ∈ X .
Then by deﬁnition E (cid:2) δ ti t | i t
= i (cid:3) =
and
E (cid:2) δ tj t | j t = j (cid:3) = 0 for all i ∈ U, j ∈ V .
We pose the following assumptions on the its variance.
Assumption 1 (Bounded conditional variance) .
We assume that the variance of δ ti t is bounded.
That is, there exists { σ 2 k } k ∈ X , such that
[Equation]
Assumption 1 allows us to provide a ﬁner characterization on the variance of stochastic gradient compared to typical variance assumption in literature.
To illustrate, recall that standard assumption in the stochastic optimization literature assumes Var( g t )
=
E (cid:107)∇ Θ f t − g t (cid:107) 2 ≤ σ 2 for some universal constant σ > 0 . Consider an extreme setting, where we have exact gradient for the sampled user-item pair, i.e., g ii t = 1 p it ∇ f ti t and g ij t = 1 p jt ∇ f tj t , then we have σ k = 0 for all k ∈ X .
In contrast, the variance of g t is still non-zero.
In general setting, we can bound the variance as shown in the following proposition.
Note that the variance lower bound arises naturally form the extreme sparsity of the stochastic gradient.
Proposition 2.1.
Given Assumption 1, we have
[Equation]
(cid:88) k
[Equation] [Equation] [Equation]
) (cid:13)(cid:13) ∇
[Equation] [Equation]
Assumption 2 (Smoothness of prediction loss) .
We assume (cid:96) ( u, v ; y ) is symmetric w.r.t.
u and v for any y ∈ {− 1 , +1 } , and there exists L >
such that (cid:13)(cid:13) ∇ 2 uu (cid:96) ( · , · ; · ) (cid:13)(cid:13) 2 ≤ L, (cid:13)(cid:13) ∇ 2 uv (cid:96) ( · , · ; · ) (cid:13)(cid:13) 2 ≤ L .
The assumption on the symmetry of (cid:96) is readily satisﬁed by almost all neural network architecture.
In essence, this assumption only requires that the parameterization of embedding vector is token agnostic.
On the other hand, the spectral upper bound on the Hessian matrix is a standard assumption in optimization literature.
T HEORETICAL R ESULTS
We ﬁrst present the convergence results of FA-SGD and standard SGD for embedding learning problem formulated in (1), and discuss the advantage that FA-SGD offers when the token distribution { p k } k ∈ X is highly imbalanced.
We further propose a variant, named CF-SGD, which can estimate frequency information in an online fashion and still provably enjoys the beneﬁts of FA-SGD.
3. C ONVERGENCE OF FA-SGD AND STANDARD SGD
Theorem 3.1 (FA-SGD) .
With Assumption 1 and 2, take learning rate policy to be
[Equation]
(7) where T denotes the total number of iterations, and α = (cid:113) ( f (Θ 0 ) − f ∗ ) / (cid:0) L (cid:80)
l ∈ X p l σ 2 l (cid:1) , we have
[Equation]
Remark 3.1 (Connection with Stochastic Block Coordinate Descent) .
Our FA-SGD shares some similarities with Stochastic Block Coordinate Descent (SBCD) (Nesterov, 2012; Dang & Lan, 2015;
Richt´arik & Tak´aˇc, 2014) applied to problem (1), in the sense that each iteration we sample certain blocks of variables ( θ i t , θ j t in our case), and only update the sampled blocks by following its stochastic gradient.
Different from SBCD, the stochastic gradient of the block variable g ti t in the FA-SGD is biased , as shown in (4).
Note that with unbiased stochastic gradient, SBCD method typically converges slower than standard SGD by a factor that can be as large as number of blocks.
As a concrete example, when the token distribution is uniform, SBCD converges slower than standard SGD by a factor of | X | , hence slower than FA-SGD by a factor of | X | from Corollary 3.1 developed later.
Recall that from Proposition 2.1, the variance of stochastic gradient is heavily inﬂuenced by the population gradient ∇ Θ f (Θ) , and can be huge whenever the population gradient is, presumably in the early phase of training.
This relationship is also supported by empirical ﬁndings in Zhang et al.
(Figure 2a), where the authors show that for BERT pretraining, the noise distribution in stochastic gradient g t is highly non-stationary, which has large variance in the beginning of the training and smaller variance at the end of training.
Since existing analysis of SGD in literature assumes a constant variance bound for the stochastic gradient, our observation in Proposition 2. requires an alternative analysis of SGD for problem (1).
To obtain the convergence rate of standard SGD in the presence of iterate-dependent variance (6), our key insight is to tailor the convergence analysis to the sparsity of the stochastic gradient for problem (1).
We show the convergence of standard SGD as the following.
Theorem 3.2 (Standard SGD) .
With Assumption 1 and 2, take learning rate policy to be η tk = min (cid:110)
14 L , α √ T (cid:111) , where T denotes the total number of iterations, and α = (cid:113) f
(Θ 0 ) − f ∗ L (cid:80) l ∈ X p 2 l σ 2 l , we have
[Equation]
Note that both FA-SGD and standard SGD attain a rate of O (1 / √ T ) .
Compared to existing rates of standard SGD (Ghadimi & Lan, 2013), we do not require constant variance bound on stochastic gradient, as we have discussed above.
Compared to existing rates of adaptive learning rate algorithms (Zhou et al., 2018a; Chen et al., 2018), both rates obtained here exhibits dimension-free property.
We emphasize here that due to the dimension-free nature of the bounds for both SGD and FA-SGD, we do not claim the proposed FA-SGD has better dependence on dimension, which is the main motivation of adaptive algorithms (Duchi et al., 2011; Kingma & Ba, 2014;
Reddi et al., 2019).
Instead, the major difference on the convergence of FA-SGD (8) and that of standard SGD (9) is that the former one is token-dependent .
Speciﬁcally, for FA-SGD, each token k ∈ X has its own convergence characterization, while all the tokens have the same convergence characterization in the standard
SGD.
We ﬁrst make a simple observation stating the equivalence of FA-SGD and standard SGD, when the token distribution { p k } k ∈ X is uniform.
Corollary 3.1 (Uniform Distribution) .
Suppose the user distribution { p i } i ∈ U and item distribution { p j } j ∈ V is the uniform distribution.
Then FA-SGD and standard SGD is equivalent to each other, in terms of both algorithmic execution and convergence rate.
3. W HEN DOES FA-SGD OUTPERFORM STANDARD SGD?
We show FA-SGD shines when the token distribution { p k } k ∈ X , deﬁned in (2), is highly imbalanced.
Before we present detailed discussions, we make an important remark that highly imbalanced token distributions are ubiquitous in social systems, presented in the form power-law .
Examples of such distributions include the degree of individuals in the social network (Muchnik et al., 2013); the frequency of words in natural language (Zipf, 2016); citations for academic papers (Brzezinski, 2015); number of links on the internet (Albert et al., 1999).
For more discussions on power-law distributions in social and natural systems, we refer readers to Kumamoto & Kamihigashi (2018).
In Figure 1c, 1d we plot the user and item counting distribution of Movielens-1M dataset.
One could clearly see that the user and item distributions are highly imbalanced, with a small percentages of users/items taking up the majority of rating records.
We defer details on the skewness of token distributions for Criteo dataset to Appendix C.
To illustrate the comparative advantage of FA-SGD when the token distribution { p k } k ∈ X is highly skewed.
We consider two classes of distribution families with different tail properties, one with exponential tail, and one with polynomial tail.
Corollary 3.2 (Exponential Tail) .
Let U
=
{ i n } | U | n =1 , V = { j m } | V | m =1 , where i n denote the user with n -th largest frequency, and j m denote the item with the m -th largest frequency.
Suppose
[Equation]
for some τ > 0 .
Deﬁne U T as the set of users whose frequencies are within e -factor from the highest frequency: U T =
{ i n : n ≤ 1 τ } , and V T similarly as V T = { j m : m ≤ 1 τ } .
We refer to U T as the top users, and V T as the top items.
[Figure]
Then given | U | , | V | ≥ 1 τ , the proposed FA-SGD, compared to standard SGD:
(1) Obtains
the
same
rate of convergence for the top users U T and top items V T ;
(2) E (cid:13)(cid:13) ∇ f τi n (cid:13)(cid:13) 2 can converge faster by a factor of Ω { exp ( τ ( n − | U T | )) } for i n ∈ U \ U T ; (3) E (cid:13)(cid:13) ∇ f τj m (cid:13)(cid:13) 2 can converge faster by a factor of Ω { exp ( τ ( m − | V T | )) } for j m ∈ V \ V T .
We remark that | U | , | V | ≥ 1 τ is a very mild condition, as it only requires that the most infrequent user/item should have its frequency smaller than the most frequent user/item by at least a factor of e .
i.e., the non-top user/item set U \ U T , V \ V T is nonempty, This is readily satisﬁed by the token distributions in recommendation systems and natural language modeling (Celma, 2010; Zipf, 2016), where the lowest frequency is at least orders of magnitude smaller than the highest frequency.
The factor of e in deﬁning U T , V T can also be readily replaced by any constant larger than 1 .
From Corollary 3.2, we can see that FA-SGD improves signiﬁcantly over standard SGD for user/item distribution with exponential tail.
Speciﬁcally, FA-SGD achieves the same convergence rate of top users/items compared to SGD, meanwhile it signiﬁcantly improves the convergence of the non-top users/items.
Moreover, the strength of such an improvement increases exponentially as we move towards the tail users/items .
Corollary 3.3 (Polynomial Tail) .
Let U
=
{ i n } | U | n =1 , V = { j m } | V | m =1 , where i n denote the user with n -th largest frequency, and j m denote the item with the m -th largest frequency.
Suppose
[Equation]
for some ν ≥ 2 .
Deﬁne U T as the set of users whose frequencies are within 2 -factor from the highest frequency: U T =
{ i n : n − ν ≥ 1 / 16 } , and V T similarly as V T = { j m : m − ν ≥ 1 / 16 } .
We refer to
U T as the top users, and V T as the top items.
Then given | U | , | V | ≥ 16 1 /ν , the FA-SGD, compared to standard SGD: (1) Obtains the same rate of convergence, for the top users U T and top items V T
; (2) E (cid:13)(cid:13) ∇ f τi n (cid:13)(cid:13) 2 can converge faster by a factor of Ω (cid:110)(cid:16) n | U T | (cid:17) ν (cid:111) for each i n ∈
U \ U T ; (3) E (cid:13) (cid:13) ∇ f τj m (cid:13)(cid:13) 2 can converge faster by a factor of Ω (cid:110)(cid:16) m | V T | (cid:17) ν (cid:111)
for each j m ∈ V \ V T .
We remark that polynomial tail (37) is also the prototypical example of the power law distribution class for modeling social behaviors (Kumamoto & Kamihigashi, 2018).
The constant 2 in the condition ν ≥ 2 can be replaced by any constant strictly larger than 1, with slight changes to the constant factor in the statements of the corollary.
From Corollary 3.3, we can see that FA-SGD improves signiﬁcantly over standard SGD for user/item distribution with polynomial tail.
Speciﬁcally, FA-SGD achieves the same convergence rate of top users/items compared to SGD, meanwhile it signiﬁcantly improves the convergence of the non-top users/items.
Moreover, the strength of such an improvement increases in polynomial order as we move towards the tail users/items .
3. O NLINE E STIMATION OF F REQUENCY I NFORMATION
In certain application scenarios, the token distribution { p k } k ∈ X can be unknown in advance of learning.
To apply FA-SGD, one needs to employ a preprocessing step in order to estimate the token distribution to a high accuracy, and then run the algorithm with estimated token distribution.
Such a preprocessing step often requires additional human efforts and data.
To remove such an undesirable
[Equation]
[Equation] [Equation]
Output Θ preprocessing step, below we present an online variant of FA-SGD, which uses the counter of tokens collected during training to estimate the token distribution dynamically.
We show that the proposed
C ounter-based F requency-aware S tochastic G radient D escent (CF-SGD) is able to retain the beneﬁts of FA-SGD despite unknown token distribution.
Theorem 3.3 (Counter-based FA-SGD) .
In addition to Assumption 1 and 2, suppose (cid:107)∇ f ( · ) (cid:107) ≤ G .
Take counter-based learning rate policy in Algorithm 2 to be
[Equation]
) where T denotes the total number of iterations, α = (cid:113) M f / (cid:0) L (cid:80)
l ∈ X p l σ 2 l (cid:1) and M f = f (Θ 0 ) − f ∗ + (cid:80)
k ∈ X p k
σ 2 k
/L , we have
[Equation]
for [Equation] ≥
[Equation]
We believe the assumption on gradient bound (cid:107)∇ f ( · ) (cid:107) ≤ G is not strictly necessary and can be removed with more reﬁned analysis.
Nevertheless, the requirement on T only logarithmically depends on the gradient bound G .
In addition, we highlight that the convergence characterization in Theorem 3.3 is still token-dependent.
Speciﬁcally, we can show that despite not knowing token distribution beforehand, CF-SGD can gain the same advantages that FA-SGD enjoys over SGD.
Corollary 3.4 (Exponential Tail) .
Suppose we have the same set of conditions given in Corollary 3.2, and σ/ (cid:112) L ( f (Θ 0 ) − f ∗ ) ≤ 1 .
Deﬁne U T as the set of users whose frequencies are within efactor from the highest frequency: U T
= { i n : n ≤ 1 τ } , and V T similarly as V T = { j m : m ≤ 1 τ } .
We refer to U T as the top users, and V T as the top items.
Then given | U | , | V | ≥ 1 τ , the proposed CF-SGD, compared to standard SGD:
(1) Obtains
the same rate of convergence, for
the top users U T and top items V T
; (2) E (cid:13)(cid:13) ∇ f τi n (cid:13)(cid:13) 2 can converge faster by a factor of Ω { exp ( τ ( n − | U T | )) } for i n ∈ U \ U T ; (3) E (cid:13)(cid:13) ∇ f τj m (cid:13)(cid:13) 2 can converge faster by a factor of Ω { exp ( τ ( m − | V T | )) } for j m ∈ V \ V T .
Corollary 3.5 (Polynomial Tail) .
Suppose we have the same set of conditions given in Corollary 3.3, and σ/ (cid:112) L ( f (Θ 0 ) − f ∗ ) ≤ 1 .
Deﬁne U T as the set of users whose frequencies are within 2 -factor from the highest frequency: U T =
{ i n : n − ν ≥ 1 / 16 } , and V T similarly as V T = { j m : m − ν ≥ 1 / 16 } .
We refer to U T as the top users, and V T as the top items.
Then given | U | , | V | ≥ 16 1 /ν , the FA-SGD, compared to standard SGD:
Obtains the same rate of convergence, for the top users U T and top items V T ;
(2) E (cid:13)(cid:13) ∇ f τi n (cid:13)(cid:13) 2 can converge faster by a factor of Ω (cid:110)(cid:16) n | U T | (cid:17) ν (cid:111) for each i n ∈
U \ U T ; (3) E (cid:13)(cid:13)
∇ f τj m (cid:13)(cid:13) 2 can converge faster by a factor of Ω (cid:110)(cid:16) m | V T | (cid:17) ν (cid:111) for each j m ∈ V \ V T .
The proofs of Corollary 3.4 and 3.5 follow similar lines as in the proofs of Corollary 3.2 and 3.3, which we defer to Appendix D
E XPERIMENTS
We conduct extensive experiments to verify the effectiveness of our proposed algorithms and our developed theories, on both publicly available benchmark recommendation datasets, and a largescale industrial recommendation system.
Additional experiments on learning Word2Vec embeddings (Mikolov et al., 2013a) are presented in Appendix B, demonstrating the general applicability of the FA/CF-SGD.
We list key elements of our experiment setup for benchmark datasets below.
• Datasets : Benchmark recommendation datasets MovieLens-1M 1 and Criteo 2 .
• Models : Factorization Machine (FM) (Rendle, 2010), and DeepFM (Guo et al., 2017).
• Metric : Training loss (cross-entropy loss), and test AUC (Area Under the ROC Curve).
Baseline algorithms : SGD, Adam (Kingma & Ba, , Adagrad ( Duchi et al., 2011).
Note that Adam and Adagrad are widely popular in training
recommendation
systems and language models. We also empirically verify that the token distributions for both Movielens-1M (Figure 1) and Criteo ( Appendix C) dataset are highly imbalanced, with most of the token distributions having a clear polynomially or exponentially decaying tail.
Since CF-SGD does not require frequency information, which is a huge practical beneﬁt compared to FA-SGD, in our experiments we mainly evaluate our proposed CF-SGD against the baseline algorithms.
To ensure a fair comparison, for each dataset and model type, we carefully tune the learning rate of each algorithm for best performance 3 .
We apply early stopping and stop training whenever the validation AUC do not increase for 2 consecutive epochs, which is widely adopted in practice (Tak ´ acs et al., 2009; Dacrema et al., 2021).
[Figure]
Movielens-1M: We can observe from Figure 2 that for FM and DeepFM model: (1) SGD yields the slowest convergence in training loss and AUC.
The proposed CF-SGD yields signiﬁcantly faster convergence than SGD for training loss.
In addition, CF-SGD converges even faster than the adaptive learning algorithms in the early stage of training; (3) All the algorithms eventually reaches peak
AUC around 81 . 0% , while CF-SGD attains the peak AUC much faster than baseline algorithms.
These empirical observations help us conﬁrm the effectiveness of the proposed CF-SGD algorithm.
We further make an empirical observation that draws a close connection between adaptive algorithms and CF-SGD.
We plot the second-order gradient moment maintained by Adagrad and Adam against the estimated frequency maintained by CF-SGD.
Surprisingly, the second-order gradient moment quickly develops a close-to linear relationship with the frequency information accumulated by CFSGD (Figure 3a,3b) .
This observation suggests that Adagrad and Adam are exploiting frequency information implicitly to a large extent.
https Criteo: We observe qualitative behavior of CF-SGD similar to Movielens-1M dataset, as can be seen in Figure 3c,3d, 4a,4b.
[Figure] [Figure]
[Table]
Industrial Recommendation System: We train an ultra-large industrial recommendation model with the proposed CF-SGD.
The training data contains 10 days of user-item interaction records, with ∼ 2.5 billion examples per day (25 billion examples in total).
We use around 800 features, with ∼ 100 million average number of tokens per feature.
We compare CF-SGD with Adagrad, which has been carefully tuned in production usage.
For both algorithms, we use a batch size of 64k and do one-pass training.
Different from benchmark academic datasets, we use Normalized Entropy (NE) as the evaluating metric (He et al., 2014) (smaller is better), which is the cross-entropy loss normalized by the entropy of background click through rate.
Note that due to numerous iterations of the production model, any relative improvement ∼ 0 .
02% is considered to be signiﬁcant.
In Figure 4c, 4d we compare the training NE curve CF-SGD and Adagrad, we can see that CF-SGD shows faster convergence than Adagrad (see NE difference % in Figure 4d).
From Table 1 we can observe that CF-SGD also improves over Adagrad during the serving phase.
Memory Efﬁciency: On top of the above empirical evidences showing that CF-SGD learns fast – faster than standard SGD, and comparable (if not better) to adaptive algorithms, we further highlight that CF-SGD learns cheap .
Speciﬁcally, adaptive algorithms require additional memory to store history information for each parameter.
For an embedding table of size N × d ( N tokens, d being embedding dimension), the memory needed is at least 3 N × d for Adam (ﬁrst/second-order gradient moment), and 2 N × d for Adagrad (second-order gradient moment). In sharp contrast, CF-SGD only requires N additional memory, for storing the estimated frequency. Since the choice of typical embedding dimension d exceeds 64 (Yin & Shen, 2018), adaptive algorithms require memory at least twice the size of the embedding table, while CF-SGD requires negligible memory overhead. Note the industrial recommendation model in our experiments has a size over multiple terabytes, with above 95% of consumed by embedding tables.
Doubling the memory footprint by using standard
Adam/Adagrad is infeasible in terms of both engineering and environmental concern.
C ONCLUSION
We propose (Counter-based) Frequency-aware SGD for embedding learning problems, which adopts frequency-dependent learning rate schedule for each token.
We demonstrate provable beneﬁts that
FA/CF-SGD enjoy over standard SGD for imbalanced token distributions, with extensive experiments supporting our theoretical ﬁndings.
Our empirical ﬁndings also suggest that adaptive algorithms can implicitly exploit frequency information and hence share close connections with the proposed algorithms, this connection might be helpful in the direct analysis of adaptive algorithms for embedding learning problems, which we leave as a future direction.