G IVENS C OORDINATE D ESCENT M ETHODS FOR R OTATION M ATRIX L EARNING IN T RAINABLE E MBEDDING I NDEXES
A BSTRACT
Product quantization (PQ) coupled with a space rotation, is widely used in modern approximate nearest neighbor (ANN) search systems to signiﬁcantly compress the disk storage for embeddings and speed up the inner product computation.
Existing rotation learning methods, however, minimize quantization distortion for ﬁxed 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 signiﬁcantly in an end-to-end training scenario
I NTRODUCTION
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 efﬁciency 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 efﬁcient 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 efﬁciently 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. P RODUCT Q UANTIZATION (PQ) B ASED I NDEXES
Since Jegou et al. (2010) ﬁrst 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 brieﬂy summarize the main directions below.
Coarse Quantization , also referred to as inverted ﬁle (IVF), is introduced in the original work of
Jegou et al.
(2010) to ﬁrst 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 efﬁcient 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 (cid:62) ).
Note that the above parameterization of R is differentiable w.r.t.
the (cid:0) d 2 (cid:1) 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 are not easily parallelizable to modern GPUs thus inefﬁcient 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 ( n 2 ) .
By contrast, our proposed GCD methods only require O ( n ) dimension of local parameterization.
1. T RAINABLE I NDEXES
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 (cid:107) XX (cid:62) − I (cid:107) 22 .
We have not explored this latter approach here, since we aim for applicability of the method even for standalone PQ.
1. O UR C ONTRIBUTIONS
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 efﬁcient 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 ﬁxed 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.
M ETHOD 2. R EVISIT T RAINABLE PQ I NDEX
Figure 1 illustrates an overview of an embedding indexing layer inside a typical deep retrieval twotower model.
Formally, the indexing layer deﬁnes a full quantization function T :
R m × n −→ R m × 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 φ : R m × n −→ R m × 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 (cid:62) for orthonormal rotation matrix.
Formally,
[Equation]
We omit the deﬁnition of product quantization function φ here, since it is not the focus of this paper.
Interested readers can ﬁnd the full deﬁnitions in (Jegou et al., 2010) or (Zhang et al., 2021).
The ﬁnal loss function includes two parts, a retrieval loss and a quantization distortion loss, formally,
[Equation]
[Equation]
where L ret 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 [Figure] solve the case where the input vectors X are ﬁxed (Ge et al., 2013; Jegou et al., 2010).
We will present our approach in the following sections.
2. G IVENS R OTATION
Before we present the algorithm of learning a rotation matrix, let us ﬁrst present a few mathematical preliminaries that serve as foundations to our algorithm.
For simplicity of exposition, the embedding dimension n = 2 m will always be an even number.
Deﬁnition 1.
Let R n × n denote the set of n × n matrices.
An orthogonal group O ( n ) = { A ∈ R n × n :
AA (cid:62) =
I n } is the group of all distance-preserving transformations of the Euclidean space R n .
The special orthogonal group, denoted SO ( n ) , consists of all orthogonal matrices with determinant 1.
It is straight-forward to check that axioms of group theory are satisﬁed by O ( n ) and SO ( n ) (see any introductory text on Lie groups, e.g., (Hall, 2015)).
In fact they form Lie sub-groups of the group of n × n invertible matrices GL ( n ) , both of which are manifolds embedded naturally in R n 2 .
Deﬁnition 2.
A Givens rotation, denoted as R i,j ( θ ) , deﬁnes a rotational linear transformation around axes i and j , by an angle of θ ∈ [0 , 2 π ) .
Formally,
[Equation]
j I n − j where I k denotes the k × k identity matrix.
In other words, the Givens rotation matrix R ij ( θ ) can be regarded as an identity matrix with four elements at ( i, i ) , ( i, j ) , ( j, i ) , ( j, j ) replaced.
Hurwitz (1963) showed the following important result as early as in 1963.
The essential idea is similar to Gram-Schmidt process, but specialized to rotational matrices.
Lemma 1.
Every element in SO ( n ) can be written as a product (cid:81) 1 ≤ i<j ≤ n R i,j ( θ i,j ) , though the decomposition is not unique.
In our case, to iteratively learn the rotation, we want to take θ to be small, such as uniformly in ( − (cid:15), (cid:15) ) , at each step.
While most existing work considers a single Givens rotation at each step, to our best knowledge no published work has considered multiple rotations in one steps, namely with a step rotation composed of n/ 2 independent rotations:
[Equation]
where σ is a randomly chosen permutation , σ ( i ) is the value at position i , and θ i,j ’s are again independent and uniform in ( − (cid:15), (cid:15) ) .
The fact that these rotation axes are mutually disjoint ensures that their product can be computed in a parallel fashion.
Also the convergence proof of the coordinate descent algorithm relies on estimating the amount of descent in the tangent space, which in turns requires that the rotations commute with one another.
This is naturally satisﬁed by the choice of rotations above.
Lemma 2.
The n/ 2 families of Givens rotations R σ (1) ,σ (2) ( θ 1 , 2 ) , . .
, R σ ( n − 1) ,σ ( n ) ( θ n − 1 ,n ) generate a maximally commuting subgroup within SO ( n ) .
This follows from the observation that a maximal torus of SO ( n ) is precisely given by simultaneous rotations along n/ 2 pairs of mutually orthogonal 2 d -planes (see (Bröcker & Tom Dieck, 2013)
Chapter 5).
For odd n = 2 m + 1 , simply replace n/ 2 by ( n − 1) / 2 .
Note that both groups SO ( n ) and GL ( n ) are non-commutative for all n > 2 , just as matrix multiplication is in general non-commutative.
Unlike GL (2) , however, SO (2) is commutative, since it consists of the 1 -parameter family of planar rotations.
2. B LOCK C OORDINATE D ESCENT A LGORITHM
The overall objective function, as shown in Eq.
(1), needs to be optimized with respect to R and other parameters in T , formally,
[Equation]
where we omit other parameters in T since they are not the focus of this paper, though clearly optimizing those parameters needs special techniques such as gradient straight-through estimator (Bengio et al., 2013).
Note that, in earlier PQ work, R is either ﬁxed to be I n , the identity matrix (Chen et al., 2020), or separately optimized in a non-differentiable manner, e.g., using singular value decomposition (Ge et al., 2013), which is equivalent to optimizing only the second term in Eq.
(2) with respect to only R and ﬁxing all other parameters.
These earlier approaches do not apply to our case, since we would like to optimize the rotation matrix iteratively and jointly with other parameters.
The naive update rule R ← R − α ∇ R L , however, does not preserve the requirement that R ∈ SO ( n ) .
A straight-forward remedy is to project the update back to SO ( n ) , such as with SVD or the tangent space exponential map, both of which however require essentially a diagonalization procedure, which is prohibitively expensive.
Instead we take a direct approach to ensure that R stays on SO ( n ) after every iteration.
Consider the Hurwitz parameterization in Lemma 1, since θ ij ∈ R , we can treat it as Euclidean optimization, thus a step of gradient descent looks like R ← (cid:81) 1 ≤ i<j ≤ n R ij ( θ ij + α ∇ θ ij L ) .
However, full Hurwitz product requires Ω( n 2 ) sparse matrix multiplications, which is too expensive.
Furthermore the various θ ij are non-trivially correlated.
To address both, we consider selecting a subset of Givens rotation R ij at each step, according to the directional derivative as follows.
Proposition 1.
Given an objective function L , the (un-normalized) directional derivative ddθ
[Equation]
θ =
L ( XR ij ( θ )) , 1 ≤ i < j ≤ n , is given by
[Equation]
where ( . )
ij denotes the element at i -th row and j -th column, and (cid:104)· , ·(cid:105) is the Frobenius inner product.
Proof.
Consider the map h : R → R n × n , h ( θ ) =
XR ij ( θ ) .
By linearity, h (cid:48) (0) =
XR (cid:48) ij (0) .
Thus, the chain rule gives
[Equation]
We would like to group ∇L ( X ) and X together.
Using the fact that (cid:104) u, v (cid:105) =
u (cid:62) v , the right hand side can be written as
[Equation]
From Deﬁnition 2 of Givens rotation R ij ( θ ) , we see that
R
[Equation]
ij
[Equation]
=
[Equation]

[Equation]

[Equation]
j
[Equation]
It follows that (cid:104) R (cid:48) ij (0) , Y (cid:105) =
( Y (cid:62) − Y ) ij for any square matrix Y , which completes the proof by plugging in Eq.
Finally to normalize the directional derivative, it needs to be divided by (cid:107) R (cid:48) ij (0) (cid:107) = √ 2 .
Given the above directional derivatives for each R ij ( θ ) , we have a few options to select a subset of n/ 2 disjoint ( i, j ) pairs for R ij , listed below in order of increasing computational complexities, where GCD stands for Givens coordinate descent.
• GCD-R randomly chooses a bipartite matching and descents along the n/ 2 Givens coordinates, similar to stochastic block coordinate gradient descent (Beck & Tetruashvili, 2013).
This can be done efﬁciently by ﬁrst shufﬂing the set of coordinates { 1 , . . . , n } .
• GCD-G sorts the ij pairs according to the absolute value of its directional derivatives in a decreasing order, and greedily picks pairs one by one if they forms a bipartite matching.
Since this is our main proposed algorithm, we list the details in Algorithm 1.
• GCD-S ﬁnds a bipartite matching with maximum sum of edge weights, giving the steepest descent.
The fastest exact algorithm (Kolmogorov, 2009) for maximal matching in a general weighted graph, has a O ( n 3 ) running time, which is impractical for ﬁrst-order optimization.
With any of the above options to pick n/ 2 disjoint ij pairs, we can present the full algorithm in
Algorithm 2.
Moreover, it is worth mentioning that picking non-disjoint ij pairs neither guarantees theoretical convergence as shown in Section 2.5, nor works well in practice (see Section 3.1).
2. C OMPLEXITY A NALYSIS
One iteration of Givens coordinate descent, as shown in Algorithm 2, consists of three major steps: a) computing directional derivatives
A in O ( n 3 ) but parallelizable to O ( n ) ; b) selecting n/ 2 disjoint pairs in O ( n ) by random, O ( n 2 log n ) by greedy, or O ( n 3 ) by the exact non-parallelizable blossom algorithm (Kolmogorov, 2009); c) applying rotation in O (2 n 2 ) FLOPs by a matrix multiplication between R and an incremental rotation, which is a sparse matrix with 2 n nonzero elements.
Thus, the sequential computational complexities of one iteration for the three proposed algorithms are
O ( n 3 + n + 2 n 2 )
=
O ( n 3 ) for GCD-R, O ( n 3 + n 2 log n + n 2 + 2 n 2 ) =
O ( n 3 ) for GCD-G, and
O ( n 3 + n 3 + 2 n 2 )
=
O ( n 3 ) for GCD-S.
Though the time complexity for SVD and Cayley transform’s gradient computation is also O ( n 3 ) , the derivative computation in the proposed GCD algorithms can be fully parallelized under
modern
GPUs.
Thus GCD-R and GCD-G have parallel runtimes of O ( n 2 ) and O ( n 2 log n ) respectively.
In contrast, the SVD computation can not be easily parallelized for general dense matrices (Berry et al., 2005), due to its internal QR factorization step (Golub & Van Loan, 1996).
So is the Cayley transform gradient computation, since matrix inversion by solving a linear system can not be parallelized.
2. C ONVERGENCE A NALYSIS
In this section we consider gradient descent of a loss function L : SO ( n ) → R .
Note that SO ( n ) is not a convex subset of its natural embedding space R n 2 because convex combinations of two distinct orthogonal matrices are never orthogonal.
As usual, convergence results require some assumptions about the domain and the function.
Before stating them, ﬁrst we introduce some deﬁnitions, specialized to the space SO ( n ) of orthogonal matrices with determinant 1 .
For all below lemmas, theorem and corollary, see Appendix for a complete proof.
Deﬁnition 3.
The tangent space of SO ( n ) ⊂ R n 2 at the identity I n is given by the space of antisymmetric matrices so ( n ) := { g ∈ R n × n : g + g (cid:62) = 0 } .
That of GL ( n ) is simply gl ( n ) :
= R n 2 .
Deﬁnition 4.
A geodesic on SO ( n ) is a curve of the form ξ ( t ) =
G ∗ exp( tg ) for a starting point
G ∗ ∈ SO ( n ) and a direction given by some g ∈ so ( n ) ; exp is the matrix exponential.
Deﬁnition 5.
A function L : SO ( n ) → R is called geodesically convex if for any geodesic ξ ,
L ◦ ξ : [0 , 1] → R is a convex function.
We make two assumptions on L in order to prove global convergence results for the random Givens coordinate descent method (GCD-R) on SO ( n ) .
The convergence of GCD-G and GCD-S follows by simple comparison with GCD-R at each step.
Assumption 1.
The objective function L is geodesically convex on the level set of the manifold
S ( G 0 ) :
= { G ∈ SO ( n ) :
L ( G ) <
L ( G 0 ) } and has a unique global minimum at G ∗ .
Assumption 2.
(Global Lipschitz)
For some η > 0 and all G ∈ S ( G 0 ) , g ∈ so ( n ) :
[Equation]
Assumption 2 leads to the following Lemma
Lemma 3.
For learning rate α = 1 η , the descent is sufﬁciently large for each step:
[Equation]
Lemma 3.4-3.6 of (Beck & Tetruashvili, 2013) combined with the above lemma and Assumptionﬁnally yield the following Theorem:
Theorem 1.
Let { G k } k ≥ 0 be the sequence generated by the GCD-R method, and let D n :=
nπ be the diameter of SO ( n ) .
With the two assumptions above,
[Equation]
In particular, the value of L ( G k ) approaches that of the minimum, L ( G ∗ ) , at a sub-linear rate.
By verifying Assumptions 1 and 2 for individual functions F , we have
Corollary 1.
Let the joint rotational product quantization learning objective be of the form ˜
[Equation]
R,
[Equation]
u
[Equation]
(cid:33) s.t.
u i ∈ C
i ⊂ R n/D and (cid:107)C i (cid:107) =
K, where (cid:76) Di =
u
i :=
( u 1 , . . .
, u D ) stands for vector concatenation.
Consider the specialization to a convex F , such as, F ( x ) = (cid:104) w, x (cid:105) (dot product) or F ( x ) = (cid:104) w,x (cid:105) (cid:107) w (cid:107)(cid:107) x (cid:107) (cosine similarity).
If we ﬁx the value of w ∈ R n as well as the choice of quantization components u i , gradient descent applied to the map L : R (cid:55)→ ˜ F ( R, C ) converges to the global minimum, provided the latter is unique.
[Figure]
E XPERIMENTS
We evaluate the convergence performance on ﬁxed embeddings in Section 3.1, and the overall embedding retrieval performance on end-to-end trainable
embedding indexes in Section 3.2.
We implement the proposed GCD algorithm in Tensorﬂow 1.
[Figure]
and perform all experiments in a single machine with 4 Nvidia 16G V100 GPU cards.
3. F IXED E MBEDDINGS
Dataset and Setup.
To make sure the proposed GCD algorithms practically converge, we ﬁrst evaluate, for a given set of embeddings from SIFT1M (Lowe, 2004), whether the GCD algorithms can converge similarly as the original OPQ algorithm
(Ge et al., 2013), which alternates between k-means for PQ centroid learning and SVD projection for optimal rotation.
Our GCD algorithms simply replace the latter SVD step by a number (5 in this experiment) of Givens coordinate descent iterations with learning rate λ =0.0001 as shown in Algorithm 2.
Similarly, Cayley method (Helfrich et al., 2018; Casado & Martínez-Rubio, 2019) also replaces the SVD steps by a number of Cayley parameter updates.
All methods are evaluated by quantization distortion, which is the square of L distance between the original embedding and quantized embedding, averaged over all examples.
Figure 2a shows the comparative results between OPQ, Cayley method and the proposed GCD algorithms.
Speciﬁcally, the proposed GCD-G and GCD-S can converge as well as the original
OPQ algorithm, as the convergence curves are similar.
The previously widely used method, Cayley transform, does not converge as fast as the proposed GCD algorithms, potentially due to the varying scales of the transformed gradient.
It also shows the results of two other alternatives: overlapping
GCD-G and overlapping GCD-R, which do not enforce the ( i, j ) pairs to be disjoint as we discussed in Section 2.2.
The comparison shows the necessity of enforcing the disjoint selection, since otherwise
GCD-G does not converge well, though it does not affect GCD-R much.
But GCD-R in general does not work as well as GCD-G, which indicates the necessity of picking the steeper descent directions.
Moreover, we can observe that GCD-G converges similarly as GCD-S, which indicates that a greedy bipartite matching algorithm as given in Algorithm 1 may be good enough in practice.
Figures 2b and 2c further compare the convergence performance between GCD-G and OPQ, by averaging over 10 runs and varying the size of data.
We can make two important observations: a)
GCD-G converges much more stably than OPQ, as the curves show signiﬁcantly lower variances; b)
GCD-G converges better for smaller size of data, e.g., for 10% of the data.
It indicates that GCD-G potentially works better in the stochastic descent scenario where the batch size is usually small.
3. E ND - TO -E ND T RAINABLE E MBEDDING I NDEXES
Dataset and Setup.
In this section, we evaluate GCD-S and GCD-R performances in an industrial search click log data where a user input query is used to retrieve items, and two public datasets of MovieLens (Harper & Konstan, 2015) and Amazon Books (He & McAuley, 2016) where user historical behaviors are used to retrieve next behavior.
The industrial dataset contains 9,989, training examples with 1,031,583 unique queries and 1,541,673 unique items, which is subsampled about 2% from our full dataset, but enough for training a reasonably well search retrieval model.
We evaluate the effectiveness of the learned rotation matrix by quantization distortion, precision@ k
[Figure]
Industrial Dataset
MovieLens
[Table]
(p@ k ) and recall@ k (r@ k ) metrics, which are standard retrieval quality metrics.
In our experiments, we choose k =100.
Speciﬁcally, for each query, we retrieve from embedding indexes a set of top k items, which is then compared with the ground truth set of items to calculate the precision and recall.
We perform ANOVA test to evaluate the p-value for all comparisons.
We implement a two-tower retrieval model (Zhang et al., 2021) with cosine scoring, hinge loss of margin 0.1 and embedding size 512.
At the beginning, a number of 10,000 warmup steps are performed before applying the indexing layer as shown in Figure 1.
Then a number of 8,192 examples are collected to warm start the PQ centroids and rotation matrix with 200 iterations of OPQ algorithm.
The baseline method uses this initial rotation matrix and stops updating it, while the other methods continue updating it.
As shown in Figure 3, in all three datasets, the proposed GCD algorithms and Cayley method clearly outperform the baseline in terms of quantization distortion, which consequentially translates into the improvement in retrieval quality metrics, measured by precision@ 100 and recall@ 100 as shown in Table 1.
However, the Amazon Book dataset is an exception.
We suspect the dataset may have subtle distribution, thus the small quantization distortion reduction does not result in improvements on retrieval quality metrics.
Finally, we observe that the three proposed GCD algorithms consistently with theoretical analysis – GCD-R and GCD-S can be regarded as the lower bound and upper bound of GCD methods, with GCD-G in between.
The GCD-S with steepest Givens rotations picked at each step, generally shows the best performance in both quantization distortion and retrieval metrics, especially better than the widely known Cayley method in literature.
C ONCLUSION
In this paper, we have proposed a novel method called Givens coordinate descent (GCD) to iteratively learn a rotation matrix, in the context of end-to-end trainable PQ based embedding indexes.
Based on a well studied mathematical result that any orthogonal matrix can be written as a product of Givens rotations, the proposed GCD algorithm picks n/ 2 maximally commuting Givens rotations to update the learning rotation matrix by their directional derivatives.
We also provide complexity analysis and prove convergence of the proposed algorithm under convexity and other mild assumptions.
Experimental results show that the proposed method converges more stably than previous SVD based and Cayley transform method, and it signiﬁcantly improves retrieval metrics in the end-to-end training of embedding indexes.
Our algorithm is also applicable to other scenarios in any neural network training where learning a rotation matrix is potentially helpful.
A C ONVERGENCE R ESULTS
In this section we prove the results stated in Section 2.5.
Lemma 1.
For learning rate α = 1 L , the descent is sufﬁciently large for each step:
[Equation]
Proof.
Recall that each step of RGCGD is given by x k + = x k · G where G =
[Equation]
n/ 2 i =
[Equation]
g σ (2 i − 1) ,σ (2 i ) ( θ i ) , where σ ∈ S n is a uniformly random permutation.
Furthermore since the g σ (2 i − 1) ,σ (2 i ) ’s in the product are along disjoint axes, they commute with one another:
[Equation]
we can write them as g i,j ( θ ) =
exp( θ √ 2 ( e i ∧ e j )) , where the √ 2 factor accounts for the fact that (cid:107) e i ∧ e j (cid:107) 2 = 2 under Euclidean norm.
Thus let’s deﬁne s ∈ so ( n ) :
= T I SO ( n ) by
[Equation]
In fact, the θ i ’s are projections of the full gradient ∇ ˜ F ( x ) , for ˜ F ( x ) :
[Equation]
F ( x k · x
[Equation] [Equation]
i,
[Equation]
So from (4), and orthogonality of 1 √ 2 e σ (2 i − 1) ∧ e σ (2 i ) , we have
[Equation]
Now choose α = 1 L .
We easily get
[Equation]
Finally recall that s ∈ so ( n ) is a random vector with randomness given by the uniform σ ∈ S n .
By orthogonality and the deﬁnition of θ i above, we have
[Equation]
Taking expectation of both sides of (7):
[Equation]
Next we recall Lemma 3.5 from Beck & Tetruashvili (2013)
Lemma 2.
Let { A k } k ≥ 0 be a non-negative sequence of real numbers satisfying A k − A k +1 ≥ ηA 2 k , k = 0 , 1 , . . .
, and A 0 ≤ 1 mη .
for positive η and m .
Then
A
[Equation]
[Equation]
This along with Assumption 1 leads to global convergence result below (combining Lemma 3.4 and
Theorem 3.6 of Beck & Tetruashvili (2013)):
Theorem 1.
Let { x k } k ≥ 0 be the sequence generated by the RGCGD method.
With the two assumptions above,
[Equation]
In particular, the value of F ( x k ) approaches that of the minimum, F ( x ∗ ) , at a sub-linear rate.
Proof.
We can write x k = R x ∗ ( s ) , where s = γ (cid:48) (0) and γ : [0 , 1] → S ( x 0 ) ⊂
SO ( n ) is a constant speed geodesic path from x ∗ to x k .
This means that (cid:107) γ (cid:48) ( t ) (cid:107) is constant for all t .
By geodesic convexity, F ◦ γ : [0 , 1] → R is convex, therefore
[Equation]
Since γ (cid:48) (0) =
d ( x k , x ∗ ) , the geodesic distance between x k and x ∗ , by the constant speed parametrization of γ , Cauchy-Schwarz inequality then implies
[Equation]
This combined with Lemma 1 gives
[Equation]
Now consider the Taylor expansion around the global minimum x ∗ .
Since the gradient vanishes, the
Lipschitz condition (4) simpliﬁes to
[Equation]
For the last inequality, we again use γ (cid:48) (0) =
d ( x k , x ∗ ) ≤ R ( x 0 ) by deﬁnition of the latter.
So letting A k = E [ F ( x k ) − F ( x ∗ )] , η = 1 2 R 2 ( x 0 )( n − 1) L , and m = 1 η ( F ( x 0 ) − F ( x ∗ )) , Lemma 2 implies that
[Equation]
Corollary 1.
Let the joint rotational product quantization learning objective be given by
[Equation]
and (cid:107)C [Equation] (cid:107) = K
[Equation]
Consider the specialization of w ∈ R n and F ( w )( x ) = (cid:104) w, x (cid:105) (dot product) or F ( w )( x ) = (cid:104) w,x (cid:105) (cid:107) w (cid:107)(cid:107) x (cid:107) (cosine similarity).
If we ﬁx the value of w ∈ W as well as the choice of quantization components u i , gradient descent applied to the map ˜ F : R (cid:55)→ L ( w, R, C ) converges to the global minimum, provided the latter is unique.
Proof.
In the dot product case, ˜ F ( R ) =
(cid:104) w, Rx (cid:105) .
If there is a unique global minimizer, it can be given by UV ⊥ from the following singular value decomposition:
[Equation]
First we verify Assumption 1, namely that the following function is convex:
[Equation]
Indeed this follows immediately from diagonalizing A , which yields a linear combination of convex functions of t .
For assumption 2, recall we need to show ˜ F ( G 0 exp( A )) ≤ ˜ F ( G 0 )
+ (cid:104)∇ ˜ F ( G 0 ) ,
A (cid:105) + L 2 (cid:107) A (cid:107) 2 , for some L > 0 .
By differentiating (cid:104) w, G 0 exp( tA ) x (cid:105) at t = 0 , we see that
[Equation]
For (cid:107) A (cid:107) small, the last term can be bounded by e (cid:107) A (cid:107) − 1 − (cid:107) A (cid:107) , which is O ( (cid:107) A (cid:107) 2 ) .
For larger (cid:107) A (cid:107) , say > C , we can choose L large enough so that LC 2 2 > sup F .
The cosine similarity case follows immediately from the dot product case since the denominator is a constant for ﬁxed w and (cid:76) Di =1 u i , thanks to the fact that (cid:107) Rx (cid:107)
= (cid:107) x (cid:107) .
Corollary 2.
Under the two assumptions 1 and 2, the steepest Givens coordinate gradient descent (SGCGD) method also converges to its unique global minimum.
Proof.
By the choice of coordinates, each step of SGCGD is dominated by the corresponding step of
RGCGD based at the same point x ∈ SO ( n ) , in other words, the estimate (12) continues to hold.
B E MPIRICAL R UNNING T IME
Figure 4 shows the comparison of per step runtime cost between GCD methods and Cayley transform for standalone QP training, with batch size chosen to be 1 .
In Figure 4a, for a fair comparison, we only compare the GCD-R method with Cayley transform, since we have not implemented GCD-G and GCD-S fully in GPU operators yet.
Note that the GCD-G method needs a sorting operator implemented in GPU (see step 5 in Algorithm 1), which is possible but needs a customized GPU operator implemented in addition to the original TensorFlow release.
On the other hand for Cayley transform, while the gradient propagation and matrix multiplication are all computed on GPU, the matrix inversion operation unfortunately cannot be parallelized, thus effectively runs on CPU.
As expected from the complexity analysis in Section 2.4, GCD-R slows down quadratically in the embedding dimension n , whereas Cayley incurs an asymptotic O ( n 3 ) cost per step.
To further explore the empirical computation cost in a completely fair setup, we compare the methods all in
CPU.
As shown in Figure 4b, GCD-G and GCD-R run much faster than the Cayley method, with about 5 times and 50 times faster running speed, respectively, which roughly reﬂect the total number of ﬂoating point computations in each method.
[Figure]
[Figure]