Givens Coordinate Descent Methods for Rotation Matrix Learning in Trainable Embedding Indexes

Product quantization ( PQ ) coupled with a space rotation , is widely used in modern approximate nearest neighbor ( ANN ) search systems to significantly compress the disk storage for embeddings and speed up the inner product computation . Existing rotation learning methods , however , minimize quantization distortion for fixed embeddings , which are not applicable to an end-to-end training scenario where embeddings are updated constantly . In this paper , based on geometric intuitions from Lie group theory , in particular the special orthogonal group SO ( n ) , we propose a family of block Givens coordinate descent algorithms to learn rotation matrix that are provably convergent on any convex objectives . Compared to the state-of-the-art SVD method , the Givens algorithms are much more parallelizable , reducing runtime by orders of magnitude on modern GPUs , and converge more stably according to experimental studies . They further improve upon vanilla product quantization significantly in an end-to-end training scenario . 1 INTRODUCTION . Search index is the core technology to enable fast information retrieval in various modern computational systems , such as web search , e-commerce search , recommendation and advertising in the past few decades . As a traditional type of search index , inverted index ( Dean , 2009 ) , which maps terms to documents in order to retrieve documents by term matching , have been the mainstream type of search index for decades , thanks to its efficiency and straightforward interpretability . Recently , with the advent of deep learning era , embedding indexes coupled with approximate nearest neighbor ( ANN ) search algorithms , have established as a promising alternative to search index ( Zhang et al. , 2020 ) and recommendation index ( Covington et al. , 2016 ; Huang et al. , 2020 ; Li et al. , 2019 ) , in part due to its learnable representations and efficient ANN algorithms . The main idea of embedding indexes is to encode users ( or queries ) and items in a latent vector space , and represent their semantic proximity in terms of inner product or cosine similarity . Embedding indexes enjoy a few appealing characteristics : a ) the embeddings can be learned to optimize downstream retrieval task of interests , and b ) items can be efficiently retrieved within tens of milliseconds . The latter leverages decades of algorithmic innovations , including a ) maximum inner product search ( MIPS ) or approximate nearest neighbors ( ANN ) , such as locality sensitive hashing ( LSH ) ( Datar et al. , 2004 ) , hierarchical navigable small world graphs ( HNSW ) ( Malkov & Yashunin , 2020 ) , space indexing by trees ( Bentley , 1975 ; Dasgupta & Freund , 2008 ; Muja & Lowe , 2014 ; Bernhardsson , 2018 ) or graphs ( Harwood & Drummond , 2016 ; Iwasaki , 2015 ) , and b ) state-of-the-art product quantization ( PQ ) based approaches ( Jegou et al. , 2010 ; Johnson et al. , 2019 ; Guo et al. , 2020 ; Wu et al. , 2017 ) . In particular , PQ based indexes and its variants have regularly claimed top spots on public benchmarks such as GIST1M , SIFT1B ( Jegou et al. , 2010 ) and DEEP10M ( Yandex & Lempitsky , 2016 ) . Moreover , along with the open-sourced libraries Faiss ( Johnson et al. , 2019 ) and ScaNN ( Guo et al. , 2020 ) , PQ based embedding indexes are widely adopted in many industrial systems ( Huang et al. , 2020 ; Zhang et al. , 2020 ) . 1.1 PRODUCT QUANTIZATION ( PQ ) BASED INDEXES . Since Jegou et al . ( 2010 ) first introduced PQ and Asymmetric Distance Computation ( ADC ) from signal processing to ANN search problem , there have been multiple lines of research focused on improving PQ based indexes . We briefly summarize the main directions below . Coarse Quantization , also referred to as inverted file ( IVF ) , is introduced in the original work of Jegou et al . ( 2010 ) to first learn a full vector quantization ( referred to as coarse quantization ) by k-means clustering and then perform product quantization over the residual of the coarse quantization . This enables non-exhaustive search of only a subsets of the clusters , which allows ANN to retrieve billions of embeddings in tens of milliseconds . More work has been done in this direction , including Inverted Multi-Index ( IMI ) ( Lempitsky , 2012 ) ; Implementation Optimization efforts have mainly been spent on the computation of ADC , including using Hamming distance for fast pruning ( Douze et al. , 2016 ) , an efficient GPU implementation of ADC lookup ( Johnson et al. , 2019 ) , and SIMD-based computation for lower bounds of ADC ( André et al. , 2015 ) . Our proposed method is fully compatible with all these implementation optimizations , since we are solely focused on the rotation matrix learning algorithms . Rotation Matrix Learning is introduced to reduce the dependencies between PQ subspaces , since PQ works best when the different subspaces are statistically independent , which may not be true in practice . As early as 2013 , Optimized PQ ( OPQ ) ( Ge et al. , 2013 ) , ITQ ( Gong et al. , 2013 ) and Cartesian k-means ( Norouzi & Fleet , 2013 ) all propose the idea of alternating between learning ( product ) vector quantization and that of the rotation matrix ; the latter can be formulated as the classic Orthogonal Procrustes problem ( Schönemann , 1966 ) that has a closed form solution in terms of singular value decomposition ( SVD ) , which is also widely used in CV ( Levinson et al. , 2020 ) . Cayley Transform ( Cayley , 1846 ) recently inspires researcher to develop end-to-end algorithm to learn rotation matrix in various neural networks , especially unitary RNN ( Helfrich et al. , 2018 ; Casado & Martínez-Rubio , 2019 ) and PQ indexes ( Wu et al. , 2017 ) . Formally , Cayley transform parameterizes a d× d rotation matrix R by R = ( I −A ) ( I +A ) −1 , where A is a skew-symmetric matrix ( i.e. , A = −A > ) . Note that the above parameterization of R is differentiable w.r.t . the ( d 2 ) parameters of A . Thus , it allows for end-to-end training of rotation matrix R , as long as the gradient of R can be obtained . However , as we will discuss in Section 2.4 , these Cayley transform based methods can not be easily parallelizable to modern GPUs thus inefficient compared to our proposed method . Another drawback of Cayley transforms is numerical instability near orthogonal matrices with -1 eigenvalues . This has been partially addressed in ( Lezcano-Casado & Martínez-Rubio , 2019 ) using approximate matrix exponential methods . However the latter incurs higher computation cost per step , due to the higher order rational functions involved . The question of minimal Euclidean dimension to embed SO ( n ) has been shown in ( Zhou et al. , 2019 ) to be of order O ( n2 ) . By contrast , our proposed GCD methods only require O ( n ) dimension of local parameterization . 1.2 TRAINABLE INDEXES . Despite its various advantages , a notable drawback of embedding indexes is the separation between model training and index building , which results in extra index building time and reduced retrieval accuracy . Typically , for a large industrial dataset with hundreds of millions of examples , it takes hours to build the index , and recall rates drop by at least 10 % ( Zhang et al. , 2021 ) . This spurred a new trend of replacing embedding indexes with jointly learnable structural indexes . Researchers argue that product quantization based embedding indexes are suboptimal , as it can not be learned jointly with a retrieval model . To overcome this drawback , researchers propose a few learnable index structures , such as tree-based ones ( Zhu et al. , 2018 ; 2019 ) , a K-D matrix one ( Gao et al. , 2020 ) and an improved tree-based one ( Zhuo et al. , 2020 ) . Though these alternative approaches have also shown improved performance , they often require highly specialized approximate training techniques , whose complexities unfortunately hinder their wide adoptions . On the other hand , learning PQ based indexes and retrieval model in an end-to-end scenario is also challenging , in large part due to the presence of non-differential operators in PQ , e.g. , arg min . Recently , however , Zhang et al . ( 2021 ) proposes an end-to-end training method which tackles the non-differentiability problem by leveraging gradient straight-through estimator ( Bengio et al. , 2013 ) , and it shows improved performance over the separately trained embedding indexes . However , learning a rotation matrix in an end-to-end training scenario remains unaddressed . As alluded to earlier , all existing methods of learning rotation matrix are based on alternative minimization procedure with an expensive non-parallelizable SVD step in each iteration . Thus it is not applicable to various back-propagation based techniques in neural network learning . Note that the orthonormality constraint on the rotation matrix is critical to ensure pairwise distances are preserved between query and item embeddings . In end to end training , however , one could relax this condition by using an L2 regularization loss of the form ‖XX > − I‖22 . We have not explored this latter approach here , since we aim for applicability of the method even for standalone PQ . 1.3 OUR CONTRIBUTIONS . In this paper , we overcome the last hurdle to fully enable end-to-end training of PQ based embedding index with retrieval models , by leveraging mathematical studies of decomposition of orthogonal group by Hurwitz ( 1963 ) as early as 1963 . Based on intuitions about maximal tori in compact Lie groups ( Hall , 2015 ) , we propose an efficient Givens coordinate gradient descent algorithm to iteratively learn the rotation matrix – starting from an identity matrix and applying a set of maximally independent ( mutually commuting ) Givens block rotations at each iteration . We open sourced the proposed algorithm , which can be easily integrated with standard neural network training algorithms , such as Adagrad and Adam . Our contribution can be summarized as follows . • Methodologically , we change the landscape of learning rotation matrix in approximate nearest neighbor ( ANN ) search from SVD based to iterative Givens rotation based , in order to be applicable to end-to-end neural network training . • Algorithmically , we propose a family of Givens coordinate block descent algorithms with complexity analysis and convergence proof of the least effective variant , GCD-R. • Empirically , we prove that for fixed embeddings the proposed algorithm is able to converge similarly as the existing rotation matrix learning algorithms , and for end-to-end training the proposed algorithm is able to learn the rotation matrix more effectively . 2 METHOD . 2.1 REVISIT TRAINABLE PQ INDEX . Figure 1 illustrates an overview of an embedding indexing layer inside a typical deep retrieval twotower model . Formally , the indexing layer defines a full quantization function T : Rm×n −→ Rm×n that maps a batch of m input embedding X in n-dimensional space to output embeddings T ( X ) , which can be decomposed into two functions : a product quantization function φ : Rm×n −→ Rm×n and a rotation function with an orthonormal matrix R. The indexing layer multiplies the input embedding X by a rotation matrix R and the product quantized embedding φ ( XR ) by its inverse matrix R−1 , which is equal to its transpose R > for orthonormal rotation matrix . Formally , T ( X ) = φ ( XR ) R > . We omit the definition of product quantization function φ here , since it is not the focus of this paper . Interested readers can find the full definitions in ( Jegou et al. , 2010 ) or ( Zhang et al. , 2021 ) . The final loss function includes two parts , a retrieval loss and a quantization distortion loss , formally , L ( X ) = Lret ( T ( X ) ) + ( 1/m ) ‖XR− φ ( XR ) ‖2 ( 1 ) where Lret denotes any retrieval loss – typically by softmax cross entropy loss or hinge loss ( Covington et al. , 2016 ; Zhang et al. , 2020 ; Huang et al. , 2020 ) . Note that straight-foward optimization of L is intractable , since the quantization function φ is not differentiable , which , however , can still be tackled by gradient straight-through estimator ( Bengio et al. , 2013 ) . Then the only remaining unaddressed problem is how to iteratively learn the rotation matrix R , since existing methods only solve the case where the input vectors X are fixed ( Ge et al. , 2013 ; Jegou et al. , 2010 ) . We will present our approach in the following sections .

Summary: The paper address rotation matrix learning during product quantization in modern ANN embedding search systems. The main contribution is addressing rotation matrix learning via gradient descent of small rotation updates. The approach relies on the decomposition of any small rotation matrix into a product of Given rotations, so that partial derivatives can be obtained in parallel, but the product causes the computation of the rotation matrix itself to be slow, in O(n^2) matrix multiplications, hence the need to select a subset of coordinates and do coordinate descent. Experiments on product quantization show a marginal improvement in results over existing approaches OPQ and Cayley.