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 a method for estimating non-linear similarities between items using Gramian estimation. This is achieved by having two separate neural networks defined for each item to be compared, which are then combined via a dot product. The proposed innovation in this paper is to use Gramian estimation for the penalty parameter of the optimization which allows for the non-linear case. Two algorithms are proposed which allow for estimation in the stochastic / online setting. Experiments are presented which appear to show good performance on some standard benchmark tasks.