Solving Multi-Codebook Quantization in the GPU

Abstract

We focus on the problem of vector compression using multi-codebook quantization (MCQ). MCQ is a generalization of k -means where the centroids arise from the combinatorial sums of entries in multiple codebooks, and has become a critical component of large-scale, state-of-the-art approximate nearest neighbour search systems. MCQ is often addressed in an iterative manner, where learning the codebooks can be solved exactly via least-squares, but finding the optimal codes results in a large number of combinatorial NP-Hard problems. Recently, we have demonstrated that an algorithm based on stochastic local search for this problem outperforms all previous approaches. In this paper we introduce a GPU implementation of our method, which achieves a $30{\times }$ 30 × speedup over a single-threaded CPU implementation. Our code is publicly available ( https://github.com/jltmtz/local-search-quantization ).