Abstract:
Embedding learning has found widespread applications in recommendation systems and natural language modeling, among other domains. To learn quality embeddings efficiently, 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 efficiency of tokendependent learning rate remains underexplored. We show that incorporating frequency information of tokens in the embedding learning problems leads to provably efficient algorithms, and demonstrate that common adaptive algorithms implicitly exploit the frequency information to a large extent. Specifically, 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 significantly lower memory. Our results are the first to show token-dependent learning rate provably improves convergence for non-convex embedding learning problems.

heading: 1 INTRODUCTION
text: 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 significant 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 efficiently 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 significant 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 efficiency 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. Specifically, 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 efficiently, as they can schedule a higher learning rate for the infrequent tokens (Duchi et al., 2011).
Despite the appealing intuition, there is a significant 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; Défossez 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 benefits over SGD?
Our contributions. We answer the previous question by showing that token frequency information can be leveraged to design provably efficient algorithms for embedding learning. Specifically,
• We propose Frequency-aware Stochastic Gradient Descent (FA-SGD), a simple modification 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 Counter-based Frequency-aware Stochastic Gradient Descent (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. Specifically, they are able to significantly improve convergence for learning infrequent tokens, while maintaining convergence speed for frequent tokens. To the best of our knowledge, our proposed algorithms are the first 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 significantly 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.

heading: 1.1 RELATED LITERATURE
text: 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; Défossez 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 efficient algorithms for embedding learning problems. Our result also does not rely on the noise to be heavy-tailed for the convergence benefits of the proposed FA/CF-SGD to take effect.
Algorithm 1 Frequency-aware Stochastic Gradient Descent Input: Total iteration number T , token frequency {pk}k∈X , and learning rate schedule {ηtk}k∈X,t∈[T ] specified by (7). Initialize: Θ0 ∈ RN×d, sample τ ∼ Unif([T ]), for t = 0, . . . τ do
(1) Sample (it, jt) ∼ D, calculate the stochastic gradient
gtit = ∇θit `(θit , θjt ; yit,jt), g t jt = ∇θjt `(θit , θjt ; yit,jt)
(2) Update parameters θt+1it = θ t it − η t itg t it , θ t+1 i = θ t i , ∀i ∈ U, i 6= it
θt+1jt = θ t jt − η t jtg t jt , θ t+1 j = θ t j , ∀j ∈ V, j 6= jt
end for Output: Θτ
Notations: For a vector/matrix, we use ‖·‖ to denotes its `2-norm/Frobenius norm. We use ‖·‖2 to denote the spectral norm of a matrix.

heading: 2 PROBLEM SETUP
text: 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 setU , letter j to index item set V , and letter k to index the union setX . The embedding learning problem can be abstracted into the following stochastic optimization problem:
min Θ∈RN×d
f(Θ) = E(i,j)∼D [`(θi, θj ; yij)] = ∑
i∈U,j∈V D(i, j)`(θi, θj ; yij). (1)
Here (i, j) denotes the user-item pair sampled from the unknown interaction distribution D, θi, θj ∈ Rd (the i, j-th row of Θ) denotes their embedding vectors respectively, and the loss `(θi, θj ; yij) denotes the prediction loss for their interaction yij ∈ {−1,+1} (e.g., logistic loss). We further let
pi = ∑ j∈V D(i, j), ∀i ∈ U ; pj = ∑ i∈U D(i, j), ∀j ∈ V, (2)
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`(θi, θj ,W; yij). 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.
gt =  0> ∇θit `(θit , θjt ; yit,jt) >
... ∇θjt `(θit , θjt ; yit,jt) >
0>  (3) At iteration t, FA-SGD samples (it, jt) ∼ D, and obtain the sparse stochastic gradient gt defined in (3). Note that only the it-th and jt-th row of gt are non-zero. One can readily verify that E(it,jt)∼D [gt] = ∇Θf(Θt). Going forward, we will denote ∇f tk as the k-th row of gradient ∇f(Θt), and gtk as the k-th row of stochastic gradient gt. Note that we have
Ejt [ gtit |it = i ] = ∇f ti /pi, Eit [ gtjt |jt = j ] = ∇f tj/pj . (4)
We further denote δtk = 1 pk f tk − gtk for all k ∈ X . Then by definition E [ δtit |it = i ] = 0 and
E [ δtjt |jt = j ] = 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 δtit is bounded. That is, there exists {σ2k}k∈X , such that
E [ ‖δit‖ 2 |it = i ] ≤ σ2i , E [ ‖δjt‖ 2 |jt = j ] ≤ σ2j , ∀i ∈ U, j ∈ V. (5)
Assumption 1 allows us to provide a finer 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(gt) = E ‖∇Θf t − gt‖
2 ≤ σ2 for some universal constant σ > 0. Consider an extreme setting, where we have exact gradient for the sampled user-item pair, i.e., giit = 1 pit ∇f tit and g i jt = 1pjt ∇f tjt , then we have σk = 0 for all k ∈ X . In contrast, the variance of gt 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∑ k∈X (1/pk − 1) ∥∥∇f tk∥∥2 ≤ Var(gt) ≤∑ k∈X pkσ 2 k + ∑ k∈X (1/pk − 1) ∥∥∇f tk∥∥2 . (6)
Assumption 2 (Smoothness of prediction loss). We assume `(u, v; y) is symmetric w.r.t. u and v for any y ∈ {−1,+1}, and there exists L > 0 such that ∥∥∇2uu`(·, ·; ·)∥∥2 ≤ L,∥∥∇2uv`(·, ·; ·)∥∥2 ≤ L. The assumption on the symmetry of ` is readily satisfied 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.

heading: 3 THEORETICAL RESULTS
text: We first 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 {pk}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 benefits of FA-SGD.

heading: 3.1 CONVERGENCE OF FA-SGD AND STANDARD SGD
text: Theorem 3.1 (FA-SGD). With Assumption 1 and 2, take learning rate policy to be ηtk = min { 1/(4L), α/ √ Tpk } , (7)
where T denotes the total number of iterations, and α = √ (f(Θ0)− f∗) / ( L ∑ l∈X plσ 2 l ) , we have
E ‖∇fτk ‖ 2 = O L (f(Θ0)− f∗) T + √ pk √∑ l∈X plσ 2 l (f(Θ 0)− f∗)L √ T  , ∀k ∈ X. (8) 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; Richtárik & Takáč, 2014) applied to problem (1), in the sense that each iteration we sample certain blocks of variables (θit , θjt in our case), and only update the sampled blocks by following its stochastic gradient. Different from SBCD, the stochastic gradient of the block variable gtit 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 influenced 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 findings in Zhang et al. (2019) (Figure 2a), where the authors show that for BERT pretraining, the noise distribution in stochastic gradient gt 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.1 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 { 1
4L , α√ T
} , where T denotes the total number of iterations, and α = √ f(Θ0)−f∗
L ∑ l∈X p 2 l σ 2 l , we have
E ‖∇fτk ‖ 2 = O L (f(Θ0)− f∗) T + √∑ l∈X p 2 l σ 2 l (f(Θ 0)− f∗)L √ T  , ∀k ∈ X. (9) 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. Specifically, 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 first make a simple observation stating the equivalence of FA-SGD and standard SGD, when the token distribution {pk}k∈X is uniform. Corollary 3.1 (Uniform Distribution). Suppose the user distribution {pi}i∈U and item distribution {pj}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.

heading: 3.2 WHEN DOES FA-SGD OUTPERFORM STANDARD SGD?
text: We show FA-SGD shines when the token distribution {pk}k∈X , defined 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 {pk}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 = {in}|U |n=1, V = {jm} |V | m=1, where in denote the user with n-th largest frequency, and jm denote the item with the m-th largest frequency. Suppose pin ∝ exp(−τn), pjm ∝ exp(−τm), ∀n ∈ [|U |],m ∈ [|V |] (10) for some τ > 0. Define UT as the set of users whose frequencies are within e-factor from the highest frequency: UT = {in : n ≤ 1τ }, and VT similarly as VT = {jm : m ≤ 1 τ }. We refer to UT as the top users, and VT as the top items.
Then given |U |, |V | ≥ 1τ , the proposed FA-SGD, compared to standard SGD: (1) Obtains the same rate of convergence, for the top users UT and top items VT ; (2) E
∥∥∇fτin∥∥2 can converge faster by a factor of Ω {exp (τ(n− |UT |))} for in ∈ U \ UT ; (3) E
∥∥∇fτjm∥∥2 can converge faster by a factor of Ω {exp (τ(m− |VT |))} for jm ∈ V \ VT . 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 \ UT , V \ VT is nonempty, This is readily satisfied 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 defining UT , VT can also be readily replaced by any constant larger than 1.
From Corollary 3.2, we can see that FA-SGD improves significantly over standard SGD for user/item distribution with exponential tail. Specifically, FA-SGD achieves the same convergence rate of top users/items compared to SGD, meanwhile it significantly 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 = {in}|U |n=1, V = {jm} |V | m=1, where in denote the user with n-th largest frequency, and jm denote the item with the m-th largest frequency. Suppose
pin ∝ n−ν , pjm ∝ m−ν , ∀n ∈ [|U |],m ∈ [|V |] (11)
for some ν ≥ 2. Define UT as the set of users whose frequencies are within 2-factor from the highest frequency: UT = {in : n−ν ≥ 1/16}, and VT similarly as VT = {jm : m−ν ≥ 1/16}. We refer to UT as the top users, and VT as the top items.
Then given |U |, |V | ≥ 161/ν , the FA-SGD, compared to standard SGD: (1) Obtains the same rate of convergence, for the top users UT and top items VT ;
(2) E ∥∥∇fτin∥∥2 can converge faster by a factor of Ω{( n|UT |)ν} for each in ∈ U \ UT ;
(3) E ∥∥∇fτjm∥∥2 can converge faster by a factor of Ω{( m|VT |)ν} for each jm ∈ V \ VT .
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 significantly over standard SGD for user/item distribution with polynomial tail. Specifically, FA-SGD achieves the same convergence rate of top users/items compared to SGD, meanwhile it significantly 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.

heading: 3.3 ONLINE ESTIMATION OF FREQUENCY INFORMATION
text: In certain application scenarios, the token distribution {pk}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
Algorithm 2 Counter-based Frequency-aware Stochastic Gradient Descent Input: Total iteration number T . Initialize: Θ0 ∈ RN×d, counter sample τ ∼ Unif({T/2, . . . , T}). for t = 0, . . . τ do
(1) Sample (it, jt) ∼ D, calculate the stochastic gradient
gtit = ∇θit `(θit , θjt ; yit,jt), g t jt = ∇θjt `(θit , θjt ; yit,jt) (2) Compute counter-based learning rate η̂tit ( ctit ) , η̂tjt ( ctjt )
specified by (12) (3) Update parameters
θt+1it = θ t it − η̂ t itg t it , θ t+1 i = θ t i , ∀i ∈ U, i 6= it θt+1jt = θ t jt − η̂ t jtg t jt , θ t+1 j = θ t j , ∀j ∈ V, j 6= jt
(4) Update counters ct+1it = c t it + 1, c t+1 i = c t i, ∀i ∈ U, i 6= it
ct+1jt = c t jt + 1, c t+1 j = c t j , ∀j ∈ V, j 6= jt
end for 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 Counter-based Frequency-aware Stochastic Gradient Descent (CF-SGD) is able to retain the benefits of FA-SGD despite unknown token distribution. Theorem 3.3 (Counter-based FA-SGD). In addition to Assumption 1 and 2, suppose ‖∇f(·)‖ ≤ G. Take counter-based learning rate policy in Algorithm 2 to be
η̂tk(c t k) = min
{ 1/(4L), 1/ √ T p̂tk } , p̂tk = c t k/t, ∀k ∈ X, t ∈ [T ], (12)
where T denotes the total number of iterations, α = √ Mf/ ( L ∑ l∈X plσ 2 l ) and Mf = f(Θ0) −
f∗ + ∑ k∈X pkσ 2 k/L, we have
E ‖∇fτk ‖ 2 = O LMf T + √ pk √∑ l∈X plσ 2 l L (f(Θ 0)− f∗) √ T + √ pk (∑ l∈X plσ 2 l ) √ T  ,∀k ∈ X. (13)
for T ≥ max {
minl∈X 1 pl ,
2 logG−log(Mf (1/2L+α/ √ pk))
pk
} .
We believe the assumption on gradient bound ‖∇f(·)‖ ≤ G is not strictly necessary and can be removed with more refined analysis. Nevertheless, the requirement on T only logarithmically depends on the gradient boundG. In addition, we highlight that the convergence characterization in Theorem 3.3 is still token-dependent. Specifically, 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 σ/ √ L (f(Θ0)− f∗) ≤ 1. Define UT as the set of users whose frequencies are within efactor from the highest frequency: UT = {in : n ≤ 1τ }, and VT similarly as VT = {jm : m ≤ 1 τ }. We refer to UT as the top users, and VT 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 UT and top items VT ; (2) E
∥∥∇fτin∥∥2 can converge faster by a factor of Ω {exp (τ(n− |UT |))} for in ∈ U \ UT ; (3) E
∥∥∇fτjm∥∥2 can converge faster by a factor of Ω {exp (τ(m− |VT |))} for jm ∈ V \ VT . Corollary 3.5 (Polynomial Tail). Suppose we have the same set of conditions given in Corollary 3.3, and σ/ √ L (f(Θ0)− f∗) ≤ 1. Define UT as the set of users whose frequencies are within
2-factor from the highest frequency: UT = {in : n−ν ≥ 1/16}, and VT similarly as VT = {jm : m−ν ≥ 1/16}. We refer to UT as the top users, and VT as the top items.
Then given |U |, |V | ≥ 161/ν , the FA-SGD, compared to standard SGD: (1) Obtains the same rate of convergence, for the top users UT and top items VT ;
(2) E ∥∥∇fτin∥∥2 can converge faster by a factor of Ω{( n|UT |)ν} for each in ∈ U \ UT ;
(3) E ∥∥∇fτjm∥∥2 can converge faster by a factor of Ω{( m|VT |)ν} for each jm ∈ V \ VT .
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

heading: 4 EXPERIMENTS
text: 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-1M1 and Criteo2. • 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, 2014), 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 benefit 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 performance3. We apply early stopping and stop training whenever the validation AUC do not increase for 2 consecutive epochs, which is widely adopted in practice (Takács et al., 2009; Dacrema et al., 2021).
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. (2) The proposed CF-SGD yields significantly 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 confirm 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.
1https://grouplens.org/datasets/movielens/1m/ 2https://ailab.criteo.com/ressources/ 3Further details on architecture and hyper-parameter choice can be found at Appendix A.
Criteo: We observe qualitative behavior of CF-SGD similar to Movielens-1M dataset, as can be seen in Figure 3c,3d, 4a,4b.
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 significant. 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 Efficiency: 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. Specifically, 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 3N ×d for Adam (first/second-order gradient moment), and 2N × 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.

heading: 5 CONCLUSION
text: We propose (Counter-based) Frequency-aware SGD for embedding learning problems, which adopts frequency-dependent learning rate schedule for each token. We demonstrate provable benefits that FA/CF-SGD enjoy over standard SGD for imbalanced token distributions, with extensive experiments supporting our theoretical findings. Our empirical findings 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.

Review summaries:
This paper proposes a frequency-aware learning algorithm for embedding learning in recommender systems. The idea of incoporating frequency information into the learning algorithm is very interesting to me. The proposed method is very simple and easy to implement with provable benefits. 

This paper proposes a counter-based learning-rate scheduler for SGD. This algorithm is designed based on the long tail distribution in recommendations and languages. The proposed algorithm enjoys a theoretic guarantee unlike other adaptive learning rate methods, e.g., Adams. Simultaneously, the authors demonstrated that the proposed algorithm achieved comparable empirical performance in two recommendation datasets. Overview, the paper is well-written, which has presented the contributions clearly.  

This paper proposes two optimization algorithms for recommendation where the token distributions are highly imbalanced. The frequency information is integrated into the optimization algorithms for fast convergence and better performance. The proposed algorithms are easy to understand and implement, and the theoretical analysis for the bounds are provided.