SOLAR: Sparse Orthogonal Learned and Random Embeddings

1 INTRODUCTION . Embedding models have been the mainstay algorithms for several machine learning applications like Information Retrieval ( IR ) ( 8 ; 2 ) and Natural Language Processing ( NLP ) ( 21 ; 16 ; 31 ; 9 ) in the last decade . Embedding models are learned spin-offs from the low-rank approximation and Matrix Factorization techniques that dominated the space of recommendation systems prior to the emergence of Deep Learning ( DL ) . The primary purpose of these models is to project a rather simple and intuitive representation of an input to an abstract low-dimensional dense vector space . This projection enables two things : 1 ) tailoring the vectors to specific downstream applications and 2 ) pre-processing and storing documents or products as vectors , thereby making the retrieval process computationally efficient ( often matrix multiplication followed by sorting , which are conducive to modern hardware like GPUs ) . Besides the computational advantage , embedding models capture the semantic relationship between queries and products . A good example is product prediction for a service like Amazon . A user-typed query has to be matched against millions of products and the best search results have to be displayed within a fraction of a second . With naive product data , it would be impossible to figure out that products with ‘ aqua ’ in their titles are actually relevant to the query ‘ water ’ . Rather , if we can project all the products to a dense low-dimensional vector space , a query can also be projected to the same space and an inner product computation can be performed with all the product vectors ( usually a dot product ) . We can then display the products with the highest inner product . These projections can be learned to encapsulate semantic information and can be continually updated to reflect temporal changes in customer preference . To the best of our knowledge , embedding models are the most prevalent ones in the industry , particularly for product and advertisement recommendations ( Amazon ’ s - DSSM ( 23 ) , Facebook ’ s DLRM ( 22 ) ) . However , the scale of these problems has blown out of proportion in the past few years prompting research in extreme classification tasks , where the number of classes runs into several million . Consequentially , approaches like Tree-based Models ( 26 ; 15 ; 1 ) and Sparse-linear Models ( 36 ; 39 ; 38 ) have emerged as powerful alternatives . Particularly , Tree-based models are much faster to train and evaluate compared to the other methods . However , most real Information Retrieval systems have dynamically changing output classes and all the extreme classification models fail to generalize to new classes with limited training data ( e.g. , new products being added to the catalogue every day ) . This has caused the resurgence of embedding models for large scale Extreme Classification ( 5 ; 29 ; 3 ; 7 ) . Our Contributions : In this paper , we argue that sparse , high dimensional , orthogonal embeddings are superior to their dense low dimensional counterparts . In this regard , we make two interesting design choices : 1 ) We design the label embeddings ( e.g.products in the catalogue ) to be high dimensional , super-sparse , and orthogonal vectors . 2 ) We fix the label embeddings throughout the training process and learn only the input embeddings ( one-sided learning ) , unlike typical dense models , where both the input and label embeddings are learned . Since we use a combination of Sparse , Orthogonal , Learned and Random embeddings , we code-name our method SOLAR . We provide a theoretical premise for SOLAR by showing that one-sided and two-sided learning are mathematically equivalent . Our choices manifest in a five-fold advantage over prior methods : • Matrix Multiplication to Inverted-Index Lookup : Sparse high dimensional embeddings can obtain a subset of labels using a mere inverted-index ( 8 ) lookup and restrict the computation and sorting to those labels . This enhances the inference speed by a large margin . • Load-balanced Inverted Index : By forcing the label embeddings to be near-orthogonal and equally sparse ( and fixing them ) , we ensure that all buckets in an inverted index are equally filled and we sample approximately the same number of labels for each input . This omits the well-known imbalanced buckets issue where we sub-sample almost all the labels for popular inputs and end up hurting the inference speed . • Lower Embedding Memory : Dense embedding models need to hold all label embeddings in GPU memory to perform real-time inference . This is not a scalable solution with millions of labels ( which is a practical industry requirement ) . On the contrary , SOLAR needs to store only few integers indices per label which is very memory efficient with modern sparse array support on all platforms . These vectors can also be used with Locality Sensitive Hashing based indexing systems like FLASH ( 34 ) . • Zero-communication : Our unique construction of label embeddings enables distributed training over multiple GPUs with zero-communication . Hence , we can afford to train on a 1.67 M book recommendation dataset and three largest extreme classification datasets and outperform the respective baselines on all 4 of them on both precision and speed . • Learning to Hash : An Inverted-Index can be perceived as a hash table where all the output classes are hashed into a few buckets ( 18 ; 33 ) . By fixing the label buckets and learning to map the inputs to the corresponding label buckets , we are doing a ‘ partial learning to hash ’ task in the hindsight ( more on this in Appendix A ) . 2 RELATED WORK . SNRM : While there have been a plethora of dense embedding models , there is only one prior work called SNRM ( Standalone Neural Ranking Model ) ( 40 ) that trains sparse embeddings for the task of suggesting documents relevant to an input query ( classic web search problem ) . In SNRM , the authors propose to learn a high dimensional output layer and sparsify it using a typical L1 or L2 regularizer . However , imposing sparsity through regularization causes lopsided inverted-index with imbalanced loads and high inference times . As we see in our experiments later , these issues lead to the poor performance of SNRM on our 1.67M product recommendation dataset . GLaS : Akin to SOLAR ’ s construction of near-orthogonal label embeddings , another recent work from Google ( 11 ) also explores the idea of enforcing orthogonality to make the labels distinguishable and thereby easier for the classifier to learn . The authors enforce it in such a way that frequently co-occurring labels have high cosine-similarity and the ones that rarely co-occur have low cosine similarity . This imposition was called a Graph Laplacian and Spreadout ( GLaS ) regularizer . However , this was done entirely in the context of dense embeddings and can not be extended to our case due to the differentiability issue . We show the comparison of SOLAR against dense embedding models with and without GLaS regularizer later on in section 5.1 . Fix your Classifier : The idea of fixing label vectors was explored in ( 13 ; 24 ) . In ( 13 ) , the authors propose to initialize the last weight matrix of popular CNN architectures with Hadamard matrices and only train the preceding layer weights . With minimal loss in precision , the number of trainable parameters can be greatly reduced . However , ‘ Fix your Classifier ’ does not scale to the the huge number of labels in the tasks of our interest as the model can not be elegantly distributed across independent workers . As shown later in section 5.2 , we observe huge performance degradation too with similar network configurations as SOLAR . Sparsifying Dense Embeddings : Several prior works have proposed to project pre-trained dense embeddings to a sparse high dimensional vectors using techniques like over-complete dictionaries ( 10 ) , denoising k-sparse auto-encoders ( 28 ; 19 ) , permutation maps on unit spheres ( 4 ) . While all these approaches vindicate the superiority of sparse vectors , none of them learn end-to-end high dimensional sparse vectors and have been confined to tasks with lot fewer labels . All other embedding models ( 5 ; 29 ; 3 ; 7 ) primarily optimize a pairwise similarity-based loss function for query-label pairs , differing in the choice of projection functions . Pairwise training needs negative sampling ( 12 ) to avoid degenerate solutions and has large training times as the number of training instances ( query-label pairs , both relevant and irrelevant ) effectively blows up . SOLAR , in addition to being sparse , also solves these challenges by learning a classifier instead of a similarity based loss function , encapsulating all labels of an input at once . Since a classifier has intrinsic negative sampling , the number of effective training samples is much lower . 3 OUR METHOD : SOLAR . In this section , we describe in detail the workflow of our algorithm SOLAR . First , we will discuss the pre-processing phase where we construct random sparse label vectors ( figure 1 ) and an inverted-index of the labels ( figure 2 ) . Then , we move to the training phase where we split the label vectors into independent contiguous components and train each of them in parallel ( figure 1 ) . In the end , we show the inference procedure where we obtain the piece-wise query vector in parallel and sparsify by retaining only top buckets from each piece . We then look up the saved inverted index to retrieve and score the candidate labels to sort and predict the best ones ( figure 3 ) . Notations : N denotes the total number of labels . D is the sparse vector dimension . K is the number of non-zeros in label vectors . B = DK is the number of buckets in each component of the vector . 1 ) Pre-processing : Construction of Label Embeddings and Inverted-Index : As presented in figure 1 , let there be N labels ( N is large , in the order of a million ) . We intend to construct K-sparse ( having K non-zero indices ) high dimensional vectors for each label . As noted earlier , a large output dimension makes training a cross-entropy loss prohibitively expensive . Therefore , inspired by recent work on zero-communication Model Parallelism ( 20 ) , we partition the large dimensional vector into K subsets and train each one independently . Each subset of the partition comprises of B buckets with exactly one non-zero index . The colored blocks on the right side in figure 1 denote the non-zero indices for each label vector . To adhere to our design principle of load-balancing , for each label , we pick the non-zero index randomly in the range of B for each of the K components . To be precise , for any label , we randomly generateK integers in the range ofB . As in most of our experiments , setK = 16 andB = 30K . This makes the overall dimension D = B ×K = 480K and a sparsity ratio of 0.000533 ( 0.0533 % ) . As an example , let the generated integers be { 18189 , 8475 , 23984 , .... , 17924 , 459 } . Then the non-zero indices of the overall vector are simply B-shifted , i.e. , { 18189 , 38475 , 83984 , .... , 437924 , 450459 } . Although any random number generator would work fine , we pick our non-zero indices using sklearn ’ s murmurhash function . It is rather straightforward to see that these vectors are near-orthogonal . The expected dot-product between any two label vectors li and lj is , E ( li T ∗ lj ) = ∑ k p ( hk ( i ) = hk ( j ) ) = K B ≈ 0 . ( 1 ) B Labels 0 1 2 3 Inverted Index0 0 N-1 1 2 3 N-2 N-4 N-3 B Labels 0 1 2 3 Inverted Index1 0 N-1 2 3 N-4 N-3 1 N-2 B Labels 0 1 2 3 Inverted Indexk-2 N-1 3 N-2N-4 1 0 N-3 2 B Labels 0 1 2 3 Inverted Indexk-1 N-13 N-2 N-4 1 0 2 N-3 Figure 2 : Inverted-Index construction for the label vectors shown in figure 1 . We construct one index for each of the K chunks . Each bucket will have the same number of labels by design ( Load-Balanced ) . Figure 2 shows the toy inverted index for the label vectors shown in figure 1 . Since we train K independent models , each model is expected to predict its own ‘ buckets of high relevance ’ . Hence we maintain K separate inverted-indexes . For any input , we accumulate the candidates from each of the K inverted-indexes and take a union of them for scoring and sorting . It is noteworthy that two unrelated labels might be pooled into the same bucket . While this sounds rather jarring from a learnability perspective , it is essential for the load-balance and also to learn a positive-only association of input tokens and true-label buckets ( more on this Appendix B ) . 2 ) Training : Figure 1 also depicts the training process ( on the left side ) . In a multilabel learning problem , each input has a variable number of true labels . We lookup all the true label vectors for an input and perform an ‘ OR ’ operation over the respective sparse label vectors . Please note that at the level of sparsity we are dealing , even with zero pairwise collisions among the non-zero indices of label vectors , we still have a super-sparse representation for the resultant ‘ OR ’ vector . We partition this combined-label vector into K parts just like before and train individual classifiers ( simple feed for- ward neural networks with 1 hidden layer ) with a binary cross entropy loss function with the B dimensional few-hot vectors . Please note that these models do not communicate with each other . Since there is no overhead of parameter sharing , training can be embarrassing parallellized across multiple GPUs ( Zero-Communication Model Parallellism ) . Input Feature Hashing : Usually , naive input tokenization like bag-of-words ( BoW ) leads to a very high dimensional input . This in turn makes the first layer of the network intractable . Hence , an elegant solution for this problem is to hash the token indices to a lower dimension ( called Feature Hashing ( 35 ) ) . In our case , we use a different random seed for each of the K models and hash the input indices to a feasible range . Although we lose some input-information in each individual model ( due to feature hash collisions ) , the variation in random seed minimizes this loss when all the models are collectively taken into account . 3 ) Inference One of the key advantages of SOLAR over dense embedding models is faster inference . As mentioned earlier , the primary reason for this is the replacement of matrix-multiplication and sorting with simple lookups and aggregations . This workflow is depicted in figure 3 . Given a tokenized input , we pass it through all the K models in parallel and obtain the respective B dimensional probability vectors . We then sort these probabilities and obtain the top-m ( m varies among 50 and 100 ) buckets for each model . These m×K integers constitute the non-zero indices of the SOLAR embedding for the input query . We can query the K inverted-indexes with the respective top-m buckets for candidates . A union of all these candidates is our target set of labels . For each of these candidates , we sum up the predicted probability scores from the corresponding buckets and sort for the top results . Filtering Noisy Labels : Random initialization of embeddings will inevitably result in unrelated labels being assigned the same bucket in any of the K subsets of the partition . Hence , a lot of the candidate labels obtained from the inverted-indexes would be irrelevant to the query . However , it is unlikely that an unrelated label would appear in the top m buckets in more than one of teh K models . Hence , we can omit the labels that appear in the topm buckets less than a threshold d times . The larger d is , the more precise the retrieved candidates will be . However , setting a large d would mean fewer candidates and thereby lower recall . Later in section 5.1 ( in table 3 ) , we experiment with multiple values of d and corroborate that d = K/2 is an optimal trade-off between precision and candidate set size . Alternate Scoring Methods : It is quite customary to sum up the log of probabilities across the K models for every candidate label as it represents the log-likelihood . Another potential strategy is to sum the logit values directly ( as logits have a wider range and are more expressive ) . In our case , we chose to assign the sum of predicted probabilities for each candidate ( as sum is just a scaled version of mean ) . The rationale behind this is shown in the following analysis assuming a multi-class classification setting . For an input x , denote Pr ( y = i ) = pi ; i ∈ { 0 , 1 , 2 , ... , N } . Let Pr ( y = b|θk ) = P jb ; b ∈ { 0 , 1 , 2 , ... , B } and k ∈ { 0 , 1 , ... , K } . Since our hash-functions are random , we have P jh ( i ) = pi + ∑ k 6=i 1h ( k ) =h ( i ) pk Hence E ( P jh ( i ) ) = pi + 1 B ∑ k 6=i pk = pi + ( 1−pi ) B . After rearrangement , we get pi = B ∗ E ( P jh ( i ) ) ( B − 1 ) − 1 B − 1 This analysis shows that the original label probabilities are linearly monotonic with expected value of the respective bucket probabilities . And since expected value is proportional to the sum across all K models , summing probabilities is a principled scheme to preserving ranking . While we primarily report the precision with sum of probabilities , we compare it against the other two heuristics ( summing log-probabilities and logits ) too in our experiments ( in table 2 ) . Time-Complexity : Since our inverted indexes are load-balanced , each bucket accommodates NB labels . Hence , the top-m buckets contribute mNB candidates . The candidates retrieved from K models are bound to have some overlapping labels , accompanied by some unrelated counterparts from the respective buckets ( since we randomly initialized the label vectors ) . In the worst case of zero overlaps , the total number of candidates would be KmNB . The aggregation of scores and frequencies of occurrence in top-m buckets is a by-product of the candidate selection step . We then omit the candidates with frequency < d. Finally , we sort the aggregated scores for the remaining candidates . Including the two sorting steps : 1 ) to get top-m buckets 2 ) to get top 5 labels , the total number of operations performed is B logm + KmNB + KmN B log 5 . A dense embedding model on the other hand needs NmK + N log 5 steps ( assuming dense vectors of dimension d = mK , since SOLAR also has mK non-zeros after inference ) . For the scale of N and mK we are dealing with ( N = 1M , K = 16 , m = 50 , B = 30K ) , SOLAR supposedly has B1+log 5 times faster inference . However , matrix multiplications have specialized processes with hardware acceleration unlike SOLAR and the practical gains would be in the order of 5x ( more in section 5.1 ) .

Summary: The paper studies the problem of document retrieval using embedding based models. It argues that performing near-neighbour search on a large number of dense embeddings hurts performance and accuracy. As an alternative, the paper proposes SOLAR (SPARSE ORTHOGONAL LEARNED AND RANDOM EMBEDDINGS), a model which uses high-dimensional, ultra sparse embeddings, on which near-neighbour search can be done using simple lookup operations. In SOLAR the document labels are divided into equal chunks of sparse vector, and independent models are learned for mapping the query to each chunk. Hence, SOLAR could be trained on multiple GPUs in an embarrassingly parallel way without requiring any communication between the GPUs. The paper demonstrates the effectiveness of SOLAR by comparing it against strong baselines on various recommendation and extreme classification datasets. It also provides theoretical justification for SOLAR by showing how “one-sided” learning (i.e. fixing label embedding but learning mapping from query to label) is mathematically equivalent to “two-sided” learning (i.e. jointly mapping the label and query to a common space).