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

TL;DR

This paper introduces Givens coordinate descent algorithms for learning rotation matrices in embedding indexes, enabling efficient, stable, and parallelizable end-to-end training for improved product quantization in ANN search.

Contribution

It proposes a novel family of Givens coordinate descent algorithms based on Lie group theory for rotation matrix learning, outperforming SVD in speed and stability.

Findings

Algorithms are highly parallelizable and faster on GPUs.

Givens methods converge more stably than SVD.

Significant improvements in end-to-end product quantization performance.

Abstract

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.