PyNNDescent

repository·master·Indexed 21 days ago

https://github.com/lmcinnes/pynndescent

A Python implementation of the Nearest Neighbor Descent algorithm for fast and accurate approximate nearest neighbor (ANN) search and k-neighbor-graph construction. Version 0.6.0 features the NNDescent class for search, PyNNDescentTransformer for scikit-learn pipelines, and a variety of distance functions in pynndescent.distances. It utilizes random projection trees for initialization and provides configurable trade-offs between speed and accuracy via parameters like epsilon, n_neighbors, and pruning_degree_multiplier.

Tokens
13.2K
Snippets
27
Records
61
Agent score
74%

What's inside pynndescent

  1. Overview of PyNNDescent for Approximate Nearest Neighbors

    master

    PyNNDescent is a Python implementation of Nearest Neighbor Descent for k-neighbor-graph construction and approximate nearest neighbor (ANN) search.

    Key features include:

    • High Accuracy: Targets 80%-100% accuracy rates.
    • Initialization: Uses random projection trees for initialization, which is particularly effective for metrics like Euclidean, Minkowski, angular, and cosine.
    • Graph Diversification: Performs pruning of the longest edges of any triangles in the graph to improve search quality.
    • Flexibility: Supports a wide variety of distance metrics out-of-the-box and allows for custom user-defined distance metrics while maintaining performance.
    • Scikit-learn Integration: Provides support for KNeighborTransformer as a drop-in replacement for algorithms requiring nearest neighbor computations.
  2. How PyNNDescent works

    master

    PyNNDescent implements the Nearest Neighbor Descent algorithm for $k$-neighbor-graph construction and approximate nearest neighbor search.

    Key features include:

    • Initialization: Uses random projection trees for initialization, which is particularly effective for metrics like Euclidean, Minkowski, angular, and cosine.
    • Graph Diversification: Performs pruning of the longest edges of any triangles in the graph to improve search quality.
    • Accuracy: Targets high accuracy rates, typically between 80% and 100%.
    • Integration: Integrates with Scikit-learn, providing KNeighborTransformer as a drop-in replacement for algorithms requiring nearest neighbor computations.
  3. Install PyNNDescent via conda or pip

    master

    PyNNDescent is a pure Python module with light requirements including numpy, scipy, scikit-learn >= 0.22, and numba >= 0.51.

    To install using conda:

    conda install -c conda-forge pynndescent

    To install using pip:

    pip install pynndescent
    conda install -c conda-forge pynndescent
  4. Build and query a PyNNDescent index

    master

    PyNNDescent provides a simple interface for approximate nearest neighbor search, similar to Scikit-learn's KDTrees or BallTrees. The workflow consists of two main steps: index construction and querying.

    1. Index Construction: Initialize the NNDescent object with your training data.
    2. Querying: Use the .query() method to find the $k$ nearest neighbors for a given query dataset.
    from pynndescent import NNDescent
    
    # Build a new search index on training data
    index = NNDescent(data)
    
    # Search the index for the 15 nearest neighbors of a test data set
    results = index.query(query_data, k=15)
  5. Perform approximate nearest neighbor search

    master

    Once a PyNNDescent index is built, you can perform searches for nearest neighbors. The search efficiency and accuracy are governed by several parameters:

    • epsilon: Controls the approximation accuracy. A higher epsilon allows the search to terminate earlier, making it faster but potentially less accurate.
    • k: The number of nearest neighbors to return.
    • search_tree_leaf_size / max_search_tree_depth: Parameters that control the structure of the search trees used to find initial candidates.

    The search algorithm uses a combination of tree-based candidate selection and graph-based traversal (using the neighbor graph) to find the closest points.

  6. How PyNNDescent performs approximate nearest neighbor search

    master

    PyNNDescent performs approximate nearest neighbor (ANN) searches using a neighbor graph as an index structure. The search follows a pruned breadth-first search pattern:

    1. Start: Choose a starting node in the graph (often the query point itself if it exists in the graph).
    2. Expand: Look at all nodes connected by an edge to the best untried candidate node.
    3. Pool: Add these neighbors to a potential candidate pool.
    4. Sort & Truncate: Sort the pool by closeness to the query point and keep only the top $k$ candidates.
    5. Iterate: Repeat the expansion process until all candidates in the pool have been tried.

    To balance speed and accuracy, the search can use an $\epsilon$ parameter. A smaller $\epsilon$ results in a faster search with potentially lower accuracy, while a larger $\epsilon$ allows for more backtracking around local optima, increasing accuracy at the cost of speed.

  7. Custom distance metrics for sparse data

    master
    If you implement custom distance metrics for an index built with sparse data, the metric function must be compatible with sparse data structures. This typically requires a different function signature than metrics designed for dense NumPy arrays. For implementation patterns, refer to the distance functions defined in pynndescent.sparse.py.
  8. How PyNNDescent achieves speed-accuracy trade-offs

    master

    PyNNDescent is an approximate nearest neighbor library that allows users to navigate the trade-off between search speed and result accuracy through three main levers:

    1. Index Construction (NNDescent params): By adjusting n_neighbors, diversify_prob, and pruning_degree_multiplier, you can build a graph that is either highly dense and accurate or sparse and extremely fast to traverse.
    2. Query-time Backtracking (epsilon): During a query, the algorithm can 'backtrack' to explore more paths. Increasing epsilon improves the chance of finding the true nearest neighbor at the cost of more computation.
    3. Graph Extraction: For tasks like UMAP or clustering where you need the neighbors of the training data itself, PyNNDescent provides direct access to the constructed graph via .neighbor_graph, bypassing the need for the query-optimization overhead introduced by .prepare().
  9. Configure threading and parallel search

    master

    PyNNDescent uses numba for high-performance execution. You can control the number of threads used during the NNDescent process and during subsequent searches.

    • NNDescent execution: The n_jobs parameter controls the number of threads. Setting n_jobs=-1 uses all available cores.
    • Batch queries: The parallel_batch_queries parameter determines whether the search function executes queries in parallel using Numba's prange.

    Note: The library internally manages thread counts by capturing the original Numba thread count and restoring it after the operation completes.

  10. How the NNDescent algorithm builds neighbor graphs

    master

    The NNDescent algorithm constructs an approximate $k$-neighbor graph by iteratively improving an initial imperfect graph. It uses the graph search technique to find better neighbors for the points already in the graph, effectively "pulling itself up by its own bootstraps."

    The Iterative Process:

    1. Initialization: Start with a random graph (each node connected to $k$ random nodes).
    2. Neighbor Discovery: For each node, measure the distance to the neighbors of its neighbors ("friends of friends").
    3. Update: If any "friend of a friend" is closer than the current neighbors, update the graph to include them, keeping only the $k$ closest.
    4. Convergence: Repeat the process until no further updates are made to the graph.

    Optimization Details:

    • Parallelism: The algorithm is implemented by computing all-pairs distances between the sets of neighbors and reverse-neighbors for each node, allowing the updates to run in parallel.
    • Initialization: PyNNDescent uses a small forest of Random Projection Trees to initialize the graph. This provides a much better starting point than a purely random graph, significantly reducing the number of expensive iterations needed to reach high accuracy.
  11. Initialize the search index with a tree

    master

    To enable fast approximate nearest neighbor search after the graph has been constructed, you can use tree_init. This builds a search forest (e.g., an RP forest or a hub-based tree) that can be used to find candidates for queries.

    When tree_init is enabled, the search process follows these steps:

    1. Forest Construction: An RP forest is built using n_trees, leaf_size, and max_rptree_depth.
    2. Graph-Informed Trees: If a graph is already available, PyNNDescent can build a 'hub-based' search tree (e.g., make_hub_tree for dense data or make_sparse_hub_tree for sparse data) which minimizes edge cuts and improves search quality.
    3. Data Reordering: The data and the neighbor graph are reordered according to the search tree leaf order to improve cache locality during search.
  12. Understand PyNNDescent performance characteristics

    master

    PyNNDescent is an approximate nearest neighbor (ANN) search implementation. Its performance is best understood through the trade-off between search accuracy (recall) and search speed (queries per second).

    Key performance takeaways:

    • High Accuracy Regime: PyNNDescent often outperforms or is highly competitive with state-of-the-art implementations (like HNSW or ONNG) when high recall is required.
    • Low Accuracy Regime: For applications where speed is paramount and high accuracy is not required, other implementations like ONNG or certain HNSW variants may offer higher queries-per-second.
    • Dataset Sensitivity: Performance varies significantly based on the dataset's dimensionality, size, distribution, and distance metric (Euclidean vs. Angular).