Efficient Training on Very Large Corpora via Gramian Estimation

1 INTRODUCTION . We consider the problem of learning a similarity function h : X × Y → R , that maps each pair of items , represented by their feature vectors ( x , y ) ∈ X × Y , to a real number h ( x , y ) , representing their similarity . We will refer to x and y as the left and right feature vectors , respectively . Many problems can be cast in this form : In natural language processing , x represents a context ( e.g . a bag of words ) , y represents a candidate word , and the target similarity measures the likelihood to observe y in context x ( Mikolov et al. , 2013 ; Pennington et al. , 2014 ; Levy & Goldberg , 2014 ) . In recommender systems , x represents a user query , y represents a candidate item , and the target similarity is a measure of relevance of item y to query x , e.g . a movie rating ( Agarwal & Chen , 2009 ) , or the likelihood to watch a given movie ( Hu et al. , 2008 ; Rendle , 2010 ) . Other applications include image similarity , where x and y are pixel-representations of images ( Bromley et al. , 1993 ; Chechik et al. , 2010 ; Schroff et al. , 2015 ) , and network embedding models ( Grover & Leskovec , 2016 ; Qiu et al. , 2018 ) , where x and y are nodes in a graph and the similarity is whether an edge connects them . A popular approach to learning similarity functions is to train an embedding representation of each item , such that items with high similarity are mapped to vectors that are close in the embedding space . A common property of such problems is that only a small subset of all possible pairs X ×Y is present in the training set , and those examples typically have high similarity . Training exclusively on observed examples has been demonstrated to yield poor generalization performance . Intuitively , when trained only on observed pairs , the model places the embedding of a given item close to similar items , but does not learn to place it far from dissimilar ones ( Shazeer et al. , 2016 ; Xin et al. , 2017 ) . Taking into account unobserved pairs is known to improve the embedding quality in many applications , including recommendation ( Hu et al. , 2008 ; Yu et al. , 2017 ) and word analogy tasks ( Shazeer et al. , 2016 ) . This is often achieved by adding a low-similarity prior on all pairs , which acts as a repulsive force between all embeddings . But because it involves a number of terms quadratic in the corpus size , this term is computationally intractable ( except in the linear case ) , and it is typically optimized using sampling : for each observed pair in the training set , a set of random unobserved pairs is sampled and used to compute an estimate of the repulsive term . But as the corpus size increases , the quality of the estimates deteriorates unless the sample size is increased , which limits scalability . In this paper , we address this issue by developing new methods to efficiently estimate the repulsive term , without sampling unobserved pairs . Our approach is inspired by matrix factorization models , which correspond to the special case of linear embedding functions . They are typically trained using alternating least squares ( Hu et al. , 2008 ) , or coordinate descent methods ( Bayer et al. , 2017 ) , which circumvent the computational burden of the repulsive term by writing it as a matrix-inner-product of two Gramians , and computing the left Gramian before optimizing over the right embeddings , and viceversa . Unfortunately , in non-linear embedding models , each update of the model parameters induces a simulateneous change in all embeddings , making it impractical to recompute the Gramians at each iteration . As a result , the Gramian formulation has been largely ignored in the non-linear setting , where models are instead trained using stochastic gradient methods with sampling of unobserved pairs , see Chen et al . ( 2016 ) . Vincent et al . ( 2015 ) were , to our knowledge , the first to attempt leveraging the Gramian formulation in the non-linear case . They consider a model where only one of the embedding functions is non-linear , and show that the gradient can be computed efficiently in that case . Their result is remarkable in that it allows exact gradient computation , but this unfortunately does not generalize to the case where both embedding functions are non-linear . Contributions We propose new methods that leverage the Gramian formulation in the non-linear case , and that , unlike previous approaches , are efficient even when both left and right embeddings are non-linear . Our methods operate by maintaining stochastic estimates of the Gram matrices , and using different variance reduction schemes to improve the quality of the estimates . We perform several experiments that show these methods scale far better than traditional sampling approaches on very large corpora . We start by reviewing preliminaries in Section 2 , then derive the Gramian-based methods and analyze them in Section 3 . We conduct large-scale experiments on the Wikipedia dataset in Section 4 , and provide additional experiments in the appendix . All the proofs are deferred to Appendix A . 2 PRELIMINARIES . 2.1 NOTATION AND PROBLEM FORMULATION . We consider models that consist of two embedding functions u : Rd×X → Rk and v : Rd×Y → Rk , which map a parameter vector1 θ ∈ Rd and feature vectors x , y to embeddings u ( θ , x ) , v ( θ , y ) ∈ Rk . The output of the model is the dot product2 of the embeddings hθ ( x , y ) = 〈u ( θ , x ) , v ( θ , y ) 〉 , where 〈· , ·〉 denotes the usual inner-product on Rk . Low-rank matrix factorization is a special case , in which the left and right embedding functions are linear in x and y . Figure 1 illustrates a non-linear model , in which each embedding function is given by a feed-forward neural network.3 We denote the training set by T = { ( xi , yi , si ) ∈ X ×Y ×R } i∈ { 1 , ... , n } , where xi , yi are the feature vectors and si is the target similarity for example i . To make notation more compact , we will use ui ( θ ) , vi ( θ ) as a shorthand for u ( θ , xi ) , v ( θ , yi ) , respectively . As discussed in the introduction , we also assume that we are given a low-similarity prior pij ∈ R for all pairs ( i , j ) ∈ { 1 , . . . , n } 2 . Given a differentiable scalar loss function ` : R× R→ R , the objective function is given by min θ∈Rd 1 n n∑ i=1 ` ( 〈ui ( θ ) , vi ( θ ) 〉 , si ) + λ n2 n∑ i=1 n∑ j=1 ( 〈ui ( θ ) , vj ( θ ) 〉 − pij ) 2 , ( 1 ) where the first term measures the loss on observed data , the second term penalizes deviations from the prior , and λ is a positive hyper-parameter that trades-off the two terms . To simplify the discussion , we will assume a uniform zero prior pij as in ( Hu et al. , 2008 ) , the general case is treated in Appendix B . To optimize this objective , existing methods rely on sampling to approximate the second term , and are usually referred to as negative sampling or candidate sampling , see Chen et al . ( 2016 ) ; Yu et al . ( 2017 ) for a survey . Due to the double sum in ( 1 ) , the quality of the sampling estimates degrades as the corpus size increases , which can significantly increase training times . This can be alleviated by increasing the sample size , but does not scale to very large corpora . 1In many applications , it is desirable for the two embedding functions u , v to share certain parameters , e.g . embeddings of categorical features common to left and right items ; hence , we use the same θ for both . 2This also includes cosine similarity models when the embedding functions u , v are normalized . 3We focus on models with this dot-product structure since they allow efficient retrieval : given x ∈ X , finding items y ∈ Y with high similarity to x reduces to a maximum inner-product search problem ( MIPS ) , which can be approximated efficiently ( Shrivastava & Li , 2014 ; Neyshabur & Srebro , 2015 ) . 2.2 GRAMIAN FORMULATION . A different approach to solving ( 1 ) , widely popular in matrix factorization , is to rewrite the double sum as the inner product of two Gram matrices . Let us denote by Uθ ∈ Rn×k the matrix of all left embeddings such that ui ( θ ) is the i-th row of Uθ , and similarly for Vθ ∈ Rn×k . Then denoting the matrix inner-product by 〈A , B〉 = ∑i , j AijBij , we can rewrite the double sum in ( 1 ) as : g ( θ ) : = 1 n2 n∑ i=1 n∑ j=1 ( UθV > θ ) 2 ij = 1 n2 〈 UθV > θ , UθV > θ 〉 . ( 2 ) Now , using the adjoint property of the inner product , we have 〈 UθV > θ , UθV > θ 〉 = 〈 U > θ Uθ , V > θ Vθ 〉 , and if we denote by u⊗ u the outer product of a vector u by itself , and define the Gram matrices Gu ( θ ) : = 1 n U > θ Uθ = 1 n n∑ i=1 ui ( θ ) ⊗ ui ( θ ) , Gv ( θ ) : = 1 n V > θ Vθ = 1 n n∑ i=1 vi ( θ ) ⊗ vi ( θ ) , ( 3 ) the penalty term becomes g ( θ ) = 〈Gu ( θ ) , Gv ( θ ) 〉 . ( 4 ) The Gramians are k × k PSD matrices , where k , the dimension of the embedding space , is much smaller than n – typically k is smaller than 1000 , while n can be arbitrarily large . Thus , the Gramian formulation ( 4 ) has a much lower computational complexity than the double sum formulation ( 2 ) , and this reformulation is at the core of alternating least squares and coordinate descent methods ( Hu et al. , 2008 ; Bayer et al. , 2017 ) , which operate by computing the exact Gramian for one side , and solving for the embeddings on the other . However , these methods do not apply in the non-linear setting due to the implicit dependence on θ , as a change in the model parameters simultaneously changes all embeddings on both sides , making it intractable to recompute the Gramians at each iteration , so the Gramian formulation has not been used when training non-linear models . In the next section , we show that it can in fact be leveraged in the non-linear case .

Summary: This paper proposes an efficient algorithm to learn  neural embedding models with a dot-product structure over very large corpora. The main method is to reformulate the objective function in terms of generalized Gramiam matrices, and maintain estimates of those matrices in the training process. The algorithm uses less time and achieves significantly better quality than sampling based methods.