Frequency-aware SGD for Efficient Embedding Learning with Provable Benefits

1 INTRODUCTION . 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 . 1.1 RELATED LITERATURE . 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 . 2 PROBLEM SETUP . 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 .

Summary: 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.