New Loss Functions for Fast Maximum Inner Product Search

1 INTRODUCTION . Maximum inner product search ( MIPS ) has become a popular paradigm for solving large scale classification and retrieval tasks . For example , in recommendation systems , user queries and documents are embedded into dense vector space of the same dimensionality and MIPS is used to find the most relevant documents given a user query ( Cremonesi et al. , 2010 ) . Similarly , in extreme classification tasks ( Dean et al. , 2013 ) , MIPS is used to predict the class label when a large number of classes , often on the order of millions or even billions are involved . Lately , MIPS has also been applied to training tasks such as scalable gradient computation in large output spaces ( Yen et al. , 2018 ) , efficient sampling for speeding up softmax computation ( Mussmann and Ermon , 2016 ) and sparse updates in end-to-end trainable memory systems ( Pritzel et al. , 2017 ) . To formally define Maximum Inner Product Search ( MIPS ) problem , consider a database X = { xi } i=1,2 , ... , N with N datapoints , where each datapoint xi ∈ Rd in a d-dimensional vector space . In the MIPS setup , given a query q ∈ Rd , we would like to find the datapoint x ∈ X that has the highest inner product with q , i.e. , we would like to identify x∗i : = arg max xi∈X 〈q , xi〉 . Exhaustively computing the exact inner product between q and N datapoints is often very expensive and sometimes infeasible . Several techniques have been proposed in the literature based on hashing and quantization to solve the approximate maximum inner product search problem efficiently , and the quantization based techniques have shown strong performance ( Ge et al. , 2014 ; Babenko and Lempitsky , 2014 ; Johnson et al. , 2017 ) . Quantizing each datapoint xi to x̃i not only reduces storage costs and memory bandwidth bottlenecks , but also permits efficient computation of distances . It avoids memory bandwidth intensive floating point operations through Hamming distance computation and look up table operations ( Norouzi et al. , 2014 ; Jegou et al. , 2011 ; Wu et al. , 2017 ) . In most traditional quantization works , the objective in the quantization procedures is to minimize the reconstruction error for the datapoints to be searched . In this paper , we propose a new class of loss functions in quantization to improve the performance of MIPS . Our contribution is threefold : • We derive a novel class of loss functions for quantization , which departs from regular reconstruction loss by weighting each pair of q and x based on its inner product value . We prove that such weighting leads to an effective loss function , which can be used by a wide class of quantization algorithms . • We devise algorithms for learning the codebook , as well as quantizing new datapoints , using the new loss functions . In particular , we give details for two families of quantization algorithms , product quantization and binary quantization . • We show that on large scale standard benchmark datasets , such as Glove100 , the change of objective yields a significant gain on the approximation of true inner product , as well as the retrieval performance . This paper is organized as follows . We first briefly review previous literature on quantization for Maximum Inner Product Search , as well as its links to ` 2 nearest neighbor search in Section 2 . Next , we give our main result , which is the derivation of our objective in Section 3 . Applications of the new loss functions to binary quantization and product quantization are given in Section 4 . Finally , we present the experimental results in Section 5 . 2 RELATED WORKS . There is a large body of similarity search literature on inner product and nearest neighbor search . We refer readers to ( Wang et al. , 2014 ; 2016 ) for a comprehensive survey . Some methods also transform MIPS problem into its equivalent form of ` 2 nearest neighbor using transformation such as ( Shrivastava and Li , 2014 ; Neyshabur and Srebro , 2014 ) , but in general are less successful than the ones that directly work in the original space . In general , these bodies of works can be divided into two families : ( 1 ) representing the data as quantized codes so that similarity computation becomes more efficient ( 2 ) pruning the dataset during the search so that only a subset of data points is considered . Typical works in the first family include binary quantization ( or binary hashing ) techniques ( Indyk and Motwani , 1998 ; Shrivastava and Li , 2014 ) and product quantization techniques ( Jegou et al. , 2011 ; Guo et al. , 2016 ) , although other families such as additive quantization ( Babenko and Lempitsky , 2014 ; Martinez et al. , 2016 ; 2018 ) and trenary quantization Zhu et al . ( 2016 ) also apply . There are many subsequent papers that extend these base approaches to more sophisticated codebook learning strategies , such as ( He et al. , 2013 ; Erin Liong et al. , 2015 ; Dai et al. , 2017 ) for binary quantization and Zhang et al . ( 2014 ) ; Wu et al . ( 2017 ) for product quantization . There are also lines of work that focus on learning transformations before quantization ( Gong et al. , 2013 ; Ge et al. , 2014 ) . Different from these methods which essentially minimize reconstruction error of the database points , we argue in Section 3 that reconstruction loss is suboptimal in the MIPS context , and any quantization method can potentially benefit from our proposed objective . The second family includes non-exhaustive search techniques such as tree search ( Muja and Lowe , 2014 ; Dasgupta and Freund , 2008 ) , graph search ( Malkov and Yashunin , 2016 ; Harwood and Drummond , 2016 ) , or hash bucketing ( Andoni et al. , 2015 ) in nearest neighbor search literature . There also exist variants of these for MIPS problem ( Ram and Gray , 2012 ; Shrivastava and Li , 2014 ) . Some of these approaches lead to larger memory requirement , or random access patterns due to the cost of constructing index structures in addition to storing original vectors . Thus they are usually used in combination with linear search quantization methods , in ways similar to inverted index ( Jegou et al. , 2011 ; Babenko and Lempitsky , 2012 ; Matsui et al. , 2015 ) . In addition , many researchers have devoted to high quality implementation of such libraries , including SPTAG Chen et al . ( 2018 ) , FAISS Johnson et al . ( 2017 ) , hnswlib Malkov and Yashunin ( 2016 ) etc .. We compared with ones available on ann-benchmarks in Section . 5 . 3 PROBLEM FORMULATION . Common quantization techniques focus on minimizing the reconstruction error ( sum of squared error ) when x is quantized to x̃ . It can be shown that minimizing the reconstruction errors is equivalent to minimizing the expected inner product quantization error under a mild condition on the query distribution . Indeed , consider the quantization objective of minimizing the expected total inner product quantization errors over the query distribution : Eq N∑ i=1 ‖〈q , xi〉 − 〈q , x̃i〉‖2 = Eq N∑ i=1 ‖〈q , xi − x̃i〉‖2 . ( 1 ) Under the assumption that q is isotropic , i.e. , E [ qqT ] = cI , where I is the identity matrix and c ∈ R+ , the objective function becomes N∑ i=1 Eq‖〈q , xi − x̃i〉‖2 = N∑ i=1 Eq ( xi − x̃i ) T qqT ( xi − x̃i ) = c N∑ i=1 ‖xi − x̃i‖2 Therefore , the objective becomes minimizing the reconstruction errors of the database points∑N i=1 ‖xi − x̃i‖2 , and this has been considered extensively in the literature . One key observation about the above objective function ( 1 ) is that it takes expectation over all possible combinations of datapoints x and queries q . However , it is easy to see that not all pairs of ( x , q ) are equally important . The approximation error on the pairs which have a high inner product is far more important since they are likely to be among the top ranked pairs and can greatly affect the search result , while for the the pairs whose inner product is low the approximation error matters much less . In other words , for a given datapoint x , we should quantize it with a bigger focus on its error with those queries which have high inner product with x . Following this key observation , we propose a new loss function by weighting the approximation error of the inner product based on the value of true inner product . More precisely , let w ( t ) ≥ 0 be a monotonically non-decreasing function , and consider the following inner-product weighted quantization error N∑ i=1 Eq [ w ( 〈q , xi〉 ) 〈q , xi − x̃i〉2 ] = N∑ i=1 ∫ w ( t ) Eq [ 〈q , xi − x̃i〉2|〈q , xi〉 = t ] dP ( 〈q , xi〉 ≤ t ) . ( 2 ) One common choice can be w ( t ) = I ( t ≥ T ) , in which case we care about all pairs whose inner product is greater or equal to certain threshold T , and disregard the rest of the pairs . 3.1 NEW LOSS FUNCTION . We decompose the inner-product weighted quantization errors based on the direction of the datapoints . We show that the new loss function ( 2 ) can be expressed as a weighted sum of the parallel and orthogonal components of the residual errors with respect to the raw datapoints . Formally , let r ( x , x̃ ) : = x− x̃ denote the quantization residual function . Given the datapoint x with ‖x‖ > 0 and its quantizer x̃ , we can decompose the residual error into two parts , the one parallel to x and the one orthogonal to x : r‖ ( x , x̃ ) : = 〈x− x̃ , x〉 · x ||x|| , ( 3 ) r⊥ ( x , x̃ ) : = ( x− x̃ ) − r‖ ( x , x̃ ) , ( 4 ) It ’ s easy to see that r ( x , x̃ ) = r‖ ( x , x̃ ) + r⊥ ( x , x̃ ) . Next we develop our new loss function . Without loss of generality , we assume : 1 . ‖q‖ = 1 . The norm of q does not matter to the ranking result ; 2 . ‖x‖ ≤ b . The norm of x is finite and bounded . Theorem 3.1 . Assuming the query q is uniformly distributed in d-dimensional unit sphere . Given the datapoint x and its quantizer x̃ , conditioned on the inner product 〈q , x〉 = t for some t > 0 , we have Eq [ 〈q , x− x̃〉2|〈q , x〉 = t ] = t2 ‖x‖2 ||r‖ ( x , x̃ ) ||2 + 1− t 2 ‖x‖2 d− 1 ||r⊥ ( x , x̃ ) ||2 . ( 5 ) Proof . First , we can decompose q : = q‖ + q⊥ with q‖ : = 〈q , x〉 · x||x|| and q⊥ : = q − q‖ where q‖ is parallel to x and q⊥ is orthogonal to x . Then , we have Eq [ 〈q , x− x̃〉2|〈q , x〉 = t ] = Eq [ 〈q‖ + q⊥ , r‖ ( x , x̃ ) + r⊥ ( x , x̃ ) 〉2|〈q , x〉 = t ] = Eq [ ( 〈q‖ , r‖ ( x , x̃ ) 〉+ 〈q⊥ , r⊥ ( x , x̃ ) 〉 ) 2|〈q , x〉 = t ] = Eq [ 〈q‖ , r‖ ( x , x̃ ) 〉2|〈q , x〉 = t ] + Eq [ 〈q⊥ , r⊥ ( x , x̃ ) 〉2|〈q , x〉 = t ] , ( 6 ) The last step uses the fact that Eq [ 〈q‖ , r‖ ( x , x̃ ) 〉〈q⊥ , r⊥ ( x , x̃ ) 〉|〈q , x〉 = t ] = 0 due to symmetry . The first term of ( 6 ) , Eq [ 〈q‖ , r‖ ( x , x̃ ) 〉2|〈q , x〉 = t ] = ‖r‖ ( x , x̃ ) ‖2Eq [ ‖q‖‖2|〈q , x〉 = t ] = ‖r‖‖2t2 ‖x‖2 . For the second term , since q⊥ is uniformly distributed in the ( d− 1 ) dimensional subspace orthogonal to x with the norm √ 1− t2‖x‖2 , we have Eq [ 〈q⊥ , r⊥ ( x , x̃ ) 〉 2|〈q , x〉 = t ] = 1− t2 ‖x‖2 d−1 ||r⊥ ( x , x̃ ) || 2 . Therefore , Eq [ 〈q , r ( x , x̃ ) 〉2|〈q , x〉 = t ] = t2 ‖x‖2 ||r‖ ( x , x̃ ) ||2 + 1− t 2 ‖x‖2 d− 1 ||r⊥ ( x , x̃ ) ||2 . Now we compute the inner-product weighted quantization error ( 2 ) for the case w ( t ) = I ( t ≥ T ) . One can do similar derivations for any reasonable w ( t ) . Given T > 0 and a datapoint x : Eq [ I ( 〈q , x〉 ≥ T ) 〈q , x− x̃〉2 ] = f ( T , ‖x‖ ) ||r‖ ( x , x̃ ) ||2 + g ( T , ‖x‖ ) d− 1 ||r⊥ ( x , x̃ ) ||2 . ( 7 ) where f ( T , ‖x‖ ) and g ( T , ‖x‖ ) are defined as f ( T , ‖x‖ ) : = ∫ ‖x‖ t=min ( T , ||x|| ) t2 ‖x‖2 dP ( 〈q , x〉 ≤ t ) , g ( T , ‖x‖ ) : = ∫ ‖x‖ t=min ( T , ||x|| ) ( 1− t 2 ‖x‖2 ) dP ( 〈q , x〉 ≤ t ) . It is easy to see that f ( T , ‖x‖ ) and g ( T , ‖x‖ ) are decreasing functions of T and increasing functions of ‖x‖ . N∑ i=1 Eq [ I ( 〈q , xi〉 ≥ T ) 〈q , xi − x̃i〉2 ] = N∑ i=1 f ( T , ‖xi‖ ) ||r‖ ( xi , x̃i ) ||2 + g ( T , ‖xi‖ ) d− 1 ||r⊥ ( xi , x̃i ) ||2 . ≤ N∑ i=1 f ( T , b ) ||r‖ ( xi , x̃i ) ||2 + g ( T , b ) d− 1 ||r⊥ ( xi , x̃i ) ||2 ∝ ( d− 1 ) λ ( T , b ) N∑ i=1 ||r‖ ( xi , x̃i ) ||2 + N∑ i=1 ||r⊥ ( xi , x̃i ) ||2 ( 8 ) where λ ( T , b ) = f ( T , b ) g ( T , b ) . 3.2 COMPUTING λ ( T , b ) The new loss functions ( 2 ) we proposed is upper bounded by ( 8 ) and the equality is achieved when all of ‖xi‖ = b for all xi . ( 8 ) can be viewed as the weighted sum of the parallel quantization errors and the orthogonal quantization errors , with respect to the original data points . Note that when λ ( T , b ) = 1 , ( 8 ) reduces to the traditional reconstruction loss . We can also characterize the asymptotic behavior of λ ( T , b ) and show that ( 1 ) λ ( T , b ) can be analytically computed , and ( 2 ) λ ( T , b ) → ( T/b ) 2 1− ( T/b ) 2 as the dimension d → ∞ . In practice , we can choose λ ( T , b ) empirically or through cross-validation . However , we found that λ ( T , b ) when d→∞ offers a very good estimate . We discuss the sensitivity λ ( T , b ) in Section . 7.5 of the Appendix . Theorem 3.2 . For b > 0 , T < b , we have λ ( T , b ) = ∫ arccos ( T/b ) 0 sind−2 θdθ∫ arccos ( T/b ) 0 sind θdθ − 1 . Proof . See the Section . 7.1 of Appendix . We furthermore prove that the limit of λ ( T , b ) exists and that it equals ( d−1 ) ( T/b ) 2 1− ( T/b ) 2 as d → ∞ . In Figure 2a , we plot λ with ( T/b ) = 0.2 and we can see it approaches its limit quickly as d grows . Theorem 3.3 . When T ≥ 0 , we have limd→∞ λ ( T , b ) = ( T/b ) 2 1− ( T/b ) 2 . Proof . See the Section . 7.2 of Appendix .

Summary: The paper proposes a novel quantization algorithm for maximum inner product search. Instead of imposing equal weights on all possible queries, the proposed framework considers giving larger weights on queries having larger inner product values. The authors shows that this derives a loss function having different weights on parallel and orthogonal component of residual of quantization. The paper shows the performance of the algorithm mainly by using the Glove1.2M dataset.