PaperlayerView exact revision on arXiv

Source-linked AI summary

Billion-scale similarity search with GPUs

Jeff Johnson, Matthijs Douze, Hervé Jégou

arXiv:1702.08734v1cs.CVcs.DBcs.DScs.IR

TL;DR

Billion-scale similarity search is difficult because high-dimensional data makes exhaustive or exact indexing impractical, while GPU utilization is challenging. The paper presents near-optimal GPU designs for exact and approximate nearest-neighbor search, outperforming prior approaches across mid- to large-scale tasks.

  • Problem

    High-dimensionality makes exhaustive and exact indexing impractical on billion-scale databases, while effectively exploiting heterogeneous GPU architectures remains challenging.

  • Method

    The paper develops register-based GPU k-selection and algorithmic layouts for exact and approximate k-nearest-neighbor search, including compressed representations.

  • Results

    The proposed similarity-search methods achieve near-optimal GPU performance and outperform previous approaches by large margins across mid- to large-scale nearest-neighbor tasks.

  • Takeaways & Limitations

    The approach enables exact k-means clustering and brute-force k-NN graph construction in less time than CPU-based approximate alternatives.

  • Takeaways & Limitations

    A fused k-selection layout is not used because its runtime gains are limited and some problem sizes suffer from lower parallelism and worse performance.

Abstract

from arXiv · show

Similarity search finds application in specialized database systems handling complex data such as images or videos, which are typically represented by high-dimensional features and require specific indexing structures. This paper tackles the problem of better utilizing GPUs for this task. While GPUs excel at data-parallel tasks, prior approaches are bottlenecked by algorithms that expose less parallelism, such as k-min selection, or make poor use of the memory hierarchy. We propose a design for k-selection that operates at up to 55% of theoretical peak performance, enabling a nearest neighbor implementation that is 8.5x faster than prior GPU state of the art. We apply it in different similarity search scenarios, by proposing optimized design for brute-force, approximate and compressed-domain search based on product quantization. In all these setups, we outperform the state of the art by large margins. Our implementation enables the construction of a high accuracy k-NN graph on 95 million images from the Yfcc100M dataset in 35 minutes, and of a graph connecting 1 billion vectors in less than 12 hours on 4 Maxwell Titan X GPUs. We have open-sourced our approach for the sake of comparison and reproducibility.

1. INTRODUCTION

The paper addresses billion-scale similarity search on GPUs, where high-dimensional data, dimensionality-related costs, and GPU utilization challenges make exhaustive or exact approaches impractical. It contributes GPU-focused k-selection and nearest-neighbor search designs, supported by experiments showing large gains over prior work.

  • Motivation: High-dimensional image and video representations make similarity search computationally and bandwidth intensive, while exploiting GPU resources effectively is nontrivial.The underlying processes may have high arithmetic complexity or data-bandwidth demands, and heterogeneous architectures pose additional utilization challenges.
  • Problem: Billion-scale k-NN graph construction is a flagship application, but exact search is impractical and prior methods such as NN-Descent have substantial memory overhead.The curse of dimensionality limits exhaustive and exact indexing approaches, motivating approximate search and compressed vector representations for memory-limited GPUs.
  • Approach: The paper focuses on product quantization codes because they are more effective than binary codes and better suited to GPU-efficient large-scale search.Binary codes introduce overhead for non-exhaustive search, while several more complex alternatives are difficult to implement efficiently on GPUs.
  • Contributions: The authors introduce a GPU k-selection algorithm that operates in fast register memory, supports kernel fusion, and includes a complexity analysis.They also propose a near-optimal algorithmic layout for exact and approximate k-nearest-neighbor search on GPUs.
  • Contributions: Experiments show large-margin improvements over previous work on mid- to large-scale nearest-neighbor tasks in single- and multi-GPU configurations.The comparison includes prior GPU state of the art suitable for billion-scale datasets with quantization codes.

2. PROBLEM STATEMENT

The section defines batched k-nearest-neighbor search over high-dimensional vector collections, primarily using L2 distance, and outlines exact and compressed-domain approximate search. Exact search evaluates pairwise distances and selects nearest neighbors, while IVFADC restricts computation to selected inverted lists using quantized representations.

  • Problem definition: Similarity search finds the k nearest neighbors of query vector x among collection vectors [y_i] using L2 distance.L2 distance is commonly used because it is optimized when learning several embeddings and has attractive linear algebra properties.
  • Batching: Batched k-selection processes n_q query vectors in parallel, selecting n_q × k elements and indices from separate arrays of lengths ℓ_i ≥ k.Batching provides flexibility across multiple CPU threads or GPU execution.
  • Exact search: Exact search computes the full pairwise distance matrix D, with its main bottleneck being the matrix multiplication XY^⊤ before row-wise k-selection.The first two terms of the squared distance can be precomputed in one pass over X and Y.
  • Compressed-domain search: IVFADC performs approximate, non-exhaustive search by quantizing database vectors and computing distances only for vectors in τ selected coarse-quantizer lists.The inverted file groups vectors into |C_1| lists by their coarse quantizer and scans τ lists linearly.
  • Product quantizer: Product quantization splits y into b sub-vectors, quantizes each independently, and produces b-byte codes with |C_2| = 256^b reproduction values.Each sub-quantizer typically has 256 reproduction values, so the code is stored as concatenated bytes.

3. GPU: OVERVIEW AND K-SELECTION

The section introduces GPU execution and memory features relevant to similarity search, then motivates a GPU k-selection design by highlighting limited parallelism, synchronization, and memory-movement challenges in prior methods.

  • GPU architecture: Nvidia GPUs execute 32-wide warps of CUDA threads, with blocks scheduled on streaming multiprocessors and sharing high-speed shared memory.Warp lanes share an instruction counter, while blocks execute on individual streaming multiprocessors.
  • GPU architecture: Register-resident structured data and warp shuffle operations enable warp-wide parallelism and storage beyond each lane’s limited local task.A lane-stride register array distributes successive values across neighboring lanes.
  • K-selection challenges: Existing k-selection algorithms are dominated by multiple global-memory passes, while explicit distance arrays may be too large to fit in memory.This limitation is especially relevant when distances are computed on-the-fly or stored only in small blocks.
  • K-selection challenges: For similarity search with k < 1000 or so, CPU max-heaps expose little data parallelism because serial tree updates cannot saturate SIMD units.GPU heap implementations additionally suffer from warp divergence and irregular, data-dependent memory movement.
  • K-selection challenges: GPU priority queues and novel small-k algorithms introduce further costs through small sorts, synchronization barriers, kernel launches, slower memories, and hierarchy overhead.The fgknn algorithm nevertheless motivates the use of parallel merges through its merge queue structure.

4. FAST K-SELECTION ON THE GPU

WarpSelect is a GPU k-selection design that keeps state in registers, scans the input once, and uses in-register odd-size sorting and merging to maintain the smallest values efficiently. Its warp-level organization supports k ≤1024 while avoiding cross-warp synchronization.

  • WarpSelect design: WarpSelect performs a single pass over the input, keeps all state in registers, and avoids cross-warp synchronization.The design targets the cost of scanning input once at peak memory bandwidth and can consume values from global memory or fused intermediate registers.
  • WarpSelect design: k ≤1024 is supported because the register file provides more storage than shared memory.Each warp performs k-selection for one input array, and sufficiently many arrays can provide full GPU occupancy.
  • Sorting networks: WarpSelect uses merge-odd and sort-odd, built from GPU-friendly in-register sorting networks implemented with lane-stride register arrays.The sorting-network approach exploits SIMD vector parallelism on the GPU.
  • Warp queues: 32 lanes process contiguous, coalesced input groups, while each lane maintains a thread queue and the warp maintains an ordered queue of the k smallest values seen.The requested k is rounded up when it is not a multiple of 32, and associated indices are carried with values.
  • Queue updates: When thread-queue updates violate the maintained ordering, the warp merges and sorts thread and warp queues so the warp queue retains the min-k elements.Values larger than a lane’s queue maximum can be rejected immediately; otherwise, insertion and odd-merge restore the invariants.

5. COMPUTATION LAYOUT

The computation layout combines GEMM, fused k-selection, tiling, and product-quantization lookup tables to implement exact and IVFADC search efficiently on GPUs. It also uses multi-pass selection and GPU replication or sharding to manage intermediate results, memory limits, and multi-GPU execution.

  • Distance computation: GEMM computes the −2⟨x_j, y_i⟩ term, while a fused kernel adds ∥y_i∥² and submits distances directly to register-based k-selection.The ∥x_j∥² term is deferred until after k-selection.
  • Tiling and memory: O(2ℓt_q) effective distance-matrix memory is achieved by tiling queries and running two independent tiles concurrently on different streams.Very large CPU inputs can use pinned-memory buffering to overlap CPU-to-GPU transfers with computation.
  • PQ lookup tables: 256 × d multiply-adds and n × b lookup-adds replace repeated distance computation when product-quantization lookup tables are used.Each of the b quantizers has 256 reproduction values, and the resulting codes are stored as sequential groups of b bytes per vector.
  • Multi-pass kernels: Two-pass k-selection reduces t_q × τ × max_i |I_i| partial results to t_q × f × k, then reduces them again to t_q × k final results.This design avoids the low parallelism of dedicating a single warp to selection over each query’s set of lists.
  • Fused kernel: Gathering lookup-table values and linearly scanning inverted lists dominate the fused-kernel runtime, so global-memory write-back is not the dominant contribution.Fusing scanning and k-selection can eliminate almost all intermediate results, but does not remove these memory-access costs.
  • Multi-GPU execution: Near linear speedup is obtained by replicating an index across R GPUs, while sharding distributes an index across S GPUs when it exceeds one GPU’s memory.Replication partitions queries across replicas; sharding gives every shard the full query set and requires an additional round of k-selection to merge partial results.

6. EXPERIMENTS & APPLICATIONS

Experiments show that WarpSelect and its fused GPU search designs substantially improve k-selection, exact and approximate nearest-neighbor search, clustering, and billion-scale graph construction. The approach achieves large speedups while maintaining high recall and scaling across multiple GPUs.

  • GPU k-selection: 1.62× and 2.01× faster than prior methods at ℓ=128000 for k=100 and k=1000, respectively.WarpSelect’s relative advantage over fgknn increases with larger k, while performance relative to Titan X peak declines for all implementations at larger k.
  • GPU k-selection: Register-resident state, no inter-warp synchronization or buffering, fusion, and odd-size networks distinguish WarpSelect from fgknn.The design also removes hierarchical partitioning and supports fusing k-selection into other kernels.
  • K-means: More than 2× faster than BIDMach’s GPU k-means on 8.1M MNIST8m images over 20 iterations.Both implementations use cuBLAS; the assignment stage uses exact search with k=1, implemented through parallel reduction rather than WarpSelect.
  • K-means: 52 minutes for exact k-means on 4 GPUs, versus 46 minutes for an approximate CPU method plus at least 56 minutes of preprocessing.The comparison uses 10^8 128-d vectors clustered to 85k centroids; the GPU method requires no preprocessing.
  • Exact nearest-neighbor search: 85 % of peak possible performance for exact search, while unfused execution is at least 25% slower.Full-array sorting is more than 10× slower than comparison methods, and the fused L2/k-selection kernel avoids an additional pass through the partial distance matrix.
  • Approximate nearest-neighbor search: 8.5× faster on SIFT1B while improving R@10 to 0.376 in 17.7 µs per query vector, versus 0.35 in 150 µs for the comparison method.The comparison uses m=8 bytes per vector and matched memory usage for similar accuracy.
  • Billion-scale graph construction: 35 minutes builds a Yfcc100M graph with accuracy above 0.8, while Deep1B graphs take 6 hours at lower quality or about half a day at higher quality.Using 8 Maxwell M40s improves performance sublinearly: ∼1.6× for m=20 and ∼1.7× for m=40.

7. CONCLUSION

The paper presents similarity-search algorithms that achieve near-optimal GPU performance despite the complexity of exploiting GPU arithmetic throughput and memory bandwidth. These methods make exact clustering and brute-force k-NN graph computation practical, supported by a released implementation for efficient similarity search.

  • 7. CONCLUSION: The presented similarity-search algorithms achieve near-optimal performance on GPUs.They are designed to exploit GPUs’ teraflops-scale arithmetic throughput and hundreds-of-gigabytes-per-second memory bandwidth.
  • 7. CONCLUSION: The approaches make exact k-means clustering and brute-force k-NN graph computation faster than CPU-based approximate methods.This enables applications that previously required complex approximate algorithms.
  • 7. CONCLUSION: A carefully engineered implementation is released to enable efficient similarity search on widely available GPUs.The work highlights GPUs’ potential for database applications, particularly as they become common in scientific workstations.

Appendix: Complexity analysis of WarpSelect

The appendix analyzes WarpSelect’s expected insertion and full-sort counts under random input permutations. It derives logarithmic bounds in the sequence length, with dependence on k, thread-queue length t, and warp width w.

  • Insertion sorts: N2 ≈ (k + t) log(c) = O(k log(ℓ/w)) expected insertion sorts, where c = ℓ/w.The approximation treats the thread queue as having seen all wc values and uses a first-order Taylor expansion.
  • Full sorts: For t = 1 and k > 1, ℓ > k, the single-lane full-sort count is π(ℓ, k, 1, 1) = k + k(Hℓ − Hk) or O(k log(ℓ)).For the special case t = 1 and k = 1, π(ℓ, 1, 1, 1) is Hℓ and converges to ln(ℓ) + γ.
  • Full sorts: For one lane and t > 1, k > 1, ℓ > k, the expected full-sort count is π(ℓ, k, t, 1) = O(k log(ℓ)/t).The bound follows because won ballots equal ⌊D/t⌋ over D successive min-k determinations, with k ≤ D ≤ ℓ.
  • Full sorts: For multiple lanes, the expected full-sort count satisfies π(ℓ, k, t, w) = O(wk log(ℓ/w)/t).An auxiliary no-interference quantity has this upper bound, and mutual interference can only reduce the number of ballots.
Loading 1702.08734v1…