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.

    Source: §1 · p. 1

  • Method

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

    Source: §1 · p. 2, §2 · p. 2, §4 · p. 4

  • 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.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

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

    Source: §7 · p. 10

    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.

    Source: §5.3 · p. 7

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.Source: §1 · p. 1

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.Source: §1 · p. 1

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.Source: §1 · pp. 1–2

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.Source: §1 · p. 2
  • 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.Source: §1 · p. 2

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.Source: §2 · p. 2

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.Source: §2 · p. 2

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.Source: §2 · p. 2

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.Source: §2 · p. 2

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.Source: §2 · pp. 2–3

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.Source: §3.1 · p. 3
  • 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.Source: §3.2 · p. 3

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.Source: §3.3 · p. 3
  • 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.Source: §3.3 · p. 4
  • 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.Source: §3.3 · p. 4

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.Source: §4 · p. 4, §4.2 · p. 5
  • 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.Source: §4.2 · p. 5

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.Source: §4.1 · p. 4, §4.2 · p. 5

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.Source: §4.2 · p. 5

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.Source: §4.2 · pp. 5–6

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.Source: §5.1 · p. 6

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.Source: §5.1 · p. 6

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.Source: §5.2 · p. 6

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.Source: §5.3 · p. 7

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.Source: §5.3 · p. 7

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.Source: §5.4 · p. 7

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.Source: §6.1 · p. 8
  • 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.Source: §6.1 · p. 8

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.Source: §6.2 · p. 8
  • 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.Source: §6.2 · p. 8

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.Source: §6.3 · pp. 8–9

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.Source: §6.4 · p. 9

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.Source: §6.4 · p. 9, §6.5 · p. 10

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.Source: §7 · p. 10
  • 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.Source: §7 · p. 10
  • 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.Source: §7 · p. 10

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.Source: Appendix: Complexity analysis of… · p. 12

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(ℓ) + γ.Source: Appendix: Complexity analysis of… · p. 12
  • 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 ≤ ℓ.Source: Appendix: Complexity analysis of… · p. 12
  • 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.Source: Appendix: Complexity analysis of… · p. 12

Related papers

Loading 1702.08734v1…