Givens Coordinate Descent Methods For Rotation Matrix Learning In Trainable Embedding Indexes

Jiang Yunjiang, Zhang Han, Qiu Yiming, Xiao Yun, Long Bo, Yang Wen-yun. The Tenth International Conference on Learning Representations 2022

[Paper]    
ICML Quantisation

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.

Similar Work

• Joint Learning Of Deep Retrieval Model And Product Quantization Based Embedding Index
• On The Needs For Rotations In Hypercubic Quantization Hashing
• Pairwise Rotation Hashing For High-dimensional Features
• Improving Bilayer Product Quantization For Billion-scale Approximate Nearest Neighbors In High Dimensions