MRPT (Multi-Rotation Projection Trees) for ANN

Overview

Approximate Nearest Neighbor (ANN) search is one of the most critical operations in modern information retrieval, recommendation systems, and machine learning pipelines. Given a query vector, ANN search finds the most similar vectors in a dataset without exhaustively comparing against every point—trading a small amount of accuracy for dramatic speedups. Among the family of tree-based ANN methods, Multi-Rotation Projection Trees (MRPT) offer a compelling combination of simplicity, speed, and automatic parameter tuning.

MRPT was introduced by Hyvönen et al. at the University of Helsinki. The core idea is to build an ensemble of randomized binary space partition trees, where each tree splits the data using random projections. A query traverses all trees simultaneously, and candidate neighbors are drawn from the intersection (or union) of the leaf nodes reached across the ensemble. The key innovation of MRPT is its ability to automatically select the optimal number of trees and voting threshold for a given recall target, eliminating the manual parameter tuning that plagues other ANN methods.

In a local semantic search system, MRPT provides an alternative indexing strategy for vector similarity search. While HNSW graphs dominate production ANN workloads, tree-based methods like MRPT are valuable for their simplicity, low memory overhead, and suitability for datasets that fit in memory but where index construction time must be minimized.

Core Concepts

Random Projection

Random projection is based on the Johnson–Lindenstrauss lemma, which states that a set of points in high-dimensional space can be projected into a much lower-dimensional space while approximately preserving pairwise distances. MRPT uses random projection vectors to partition the data: for each internal node of a tree, a random vector r is drawn, and each data point x is assigned to the left or right child based on whether the dot product r · x falls below or above a splitting threshold.

Ensemble of Trees

A single random projection tree is unreliable: two nearby points may end up on opposite sides of a random hyperplane. MRPT addresses this by building multiple independent trees (typically 5–100). During a query, the search traverses all trees and collects the candidate sets from each tree’s leaf node. Points that appear in many trees’ candidate sets are more likely to be true nearest neighbors.

Voting Threshold

The voting threshold v determines how many trees must agree that a candidate is close to the query. A point is included in the final candidate set only if it appears in at least v out of T trees’ leaf nodes. Setting v = 1 (union of all candidates) maximizes recall but requires checking many candidates. Setting v = T (intersection) minimizes candidates but risks missing true neighbors. MRPT’s autotuning selects the optimal v for a target recall.

Leaf Size and Tree Depth

The depth of each tree (or equivalently, the maximum leaf size) controls the granularity of partitioning. Deeper trees create smaller leaves, reducing the number of candidates per tree but increasing the risk of splitting true neighbors. Shallower trees produce larger leaves with more candidates, improving recall at the cost of more distance computations. The tree depth is typically set so that leaf nodes contain 10–100 points.

How It Works

Index Construction

Building an MRPT index proceeds in three phases:

1. Tree construction: For each of the T trees, recursively split the dataset. At each internal node, generate a random projection vector, compute the dot product of every point with this vector, find the median (or a quantile) of the projected values as the split threshold, and assign points to the left or right child. Recurse until leaf nodes contain at most l points.
2. Projection matrix storage: Store the random vectors and split thresholds for each internal node. These are needed at query time to route the query to the correct leaf. The total storage for projections is O(T × depth × d) where d is the dimensionality.
3. Autotuning (optional): Sample a set of queries from the training data, run them through the index at various voting thresholds, measure recall against brute-force results, and select the optimal (T, v) pair for the target recall.

Query Processing

To find the k nearest neighbors of a query vector q:

1. Traverse each tree from root to leaf by computing r · q at each internal node and following the appropriate child.
2. Collect the point IDs from each tree’s leaf node, creating a candidate list with vote counts.
3. Filter candidates by the voting threshold: keep only points with at least v votes.
4. Compute the exact distance from q to each remaining candidate.
5. Return the top k candidates by distance.

The tree traversal is O(T × depth × d) for the dot products. The candidate filtering and exact distance computation dominate runtime when many candidates survive the voting filter.

Defeatist Search vs. Priority Search

The basic MRPT query follows a “defeatist” strategy: at each internal node, the query goes left or right but never backtracks. This is fast but can miss neighbors near split boundaries. An enhanced strategy uses priority search, where the algorithm also explores the “wrong” child of nodes where the query’s projection value is close to the split threshold. This increases recall at the cost of examining more nodes and is similar to the spill tree approach.

Practical Applications

Vector Similarity Search in the Search Platform

A local semantic search platform uses MRPT as one of its ANN indexing strategies for searching over embedding vectors. When a user submits a semantic search query, the query text is encoded into a vector, and MRPT finds the most similar document vectors in the index. The autotuning feature is particularly valuable here: as the document collection grows, the index automatically adjusts its parameters to maintain the target recall level.

Image and Audio Retrieval

Content-based retrieval systems encode images or audio clips as feature vectors (often extracted from deep neural networks) and use ANN search to find similar items. MRPT’s low memory footprint makes it suitable for deployments where the index must fit in RAM alongside the application, such as on-device search in mobile or embedded systems.

Deduplication and Clustering

Near-duplicate detection requires finding all pairs of vectors within a distance threshold. MRPT’s ensemble approach naturally supports range queries: instead of retrieving the top k, retrieve all candidates within a fixed distance. This is useful for deduplicating documents, images, or user profiles in large datasets before running more expensive exact comparisons.

Implementation Patterns

Memory Layout for Cache Efficiency

Tree nodes are stored in a flat array using BFS (breadth-first) ordering, so that a node’s children are at indices 2i+1 and 2i+2. This avoids pointer chasing and improves cache locality during tree traversal. The projection vectors for each tree are stored contiguously in a matrix, enabling SIMD-accelerated dot products.

Batch Query Processing

When multiple queries arrive simultaneously, batch processing amortizes the overhead of loading tree structure into cache. All queries traverse the first level of all trees, then the second level, and so on. This “breadth-first query” pattern keeps each tree level hot in L2/L3 cache across all queries, significantly improving throughput compared to processing queries independently.

Incremental Index Updates

Random projection trees support incremental insertion by routing new points down existing trees and appending them to leaf nodes. When a leaf exceeds the maximum size, it can be split with a new random projection. Deletion is handled by marking points as deleted and periodically rebuilding trees that have too many deleted entries. However, frequent updates degrade the balance of trees, so periodic full rebuilds are recommended for write-heavy workloads.

Code Examples

MRPT Index in Python (NumPy)

This example implements a simplified MRPT index demonstrating the core concepts of random projection trees and ensemble voting.

import numpy as np
from typing import List, Tuple


class MRPTIndex:
    """Multi-Rotation Projection Trees for approximate nearest neighbors."""

    def __init__(self, n_trees: int = 10, max_leaf_size: int = 50):
        self.n_trees = n_trees
        self.max_leaf_size = max_leaf_size
        self.trees: List[dict] = []
        self.data: np.ndarray | None = None

    def fit(self, data: np.ndarray) -> None:
        """Build the index from a dataset of shape (n, d)."""
        self.data = data
        n, d = data.shape
        indices = np.arange(n)

        for _ in range(self.n_trees):
            tree = self._build_tree(indices, d)
            self.trees.append(tree)

    def _build_tree(self, indices: np.ndarray, dim: int) -> dict:
        """Recursively build a random projection tree."""
        if len(indices) <= self.max_leaf_size:
            return {"leaf": True, "indices": indices}

        # Generate a random projection vector.
        projection = np.random.randn(dim)
        projection /= np.linalg.norm(projection)

        # Project data points onto the random vector.
        projected = self.data[indices] @ projection

        # Split at the median.
        threshold = np.median(projected)
        left_mask = projected <= threshold
        right_mask = ~left_mask

        # Handle edge case where all points land on one side.
        if not np.any(left_mask) or not np.any(right_mask):
            return {"leaf": True, "indices": indices}

        return {
            "leaf": False,
            "projection": projection,
            "threshold": threshold,
            "left": self._build_tree(indices[left_mask], dim),
            "right": self._build_tree(indices[right_mask], dim),
        }

    def _query_tree(self, tree: dict, query: np.ndarray) -> np.ndarray:
        """Traverse a single tree to find the leaf for a query."""
        if tree["leaf"]:
            return tree["indices"]

        proj_val = query @ tree["projection"]
        if proj_val <= tree["threshold"]:
            return self._query_tree(tree["left"], query)
        else:
            return self._query_tree(tree["right"], query)

    def query(
        self, q: np.ndarray, k: int = 10, voting_threshold: int = 1
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Find the k approximate nearest neighbors.

        Args:
            q: Query vector of shape (d,).
            k: Number of neighbors to return.
            voting_threshold: Minimum number of trees that must
                              agree on a candidate.

        Returns:
            (indices, distances) of the k nearest neighbors.
        """
        # Collect votes across all trees.
        votes: dict[int, int] = {}
        for tree in self.trees:
            candidates = self._query_tree(tree, q)
            for idx in candidates:
                votes[idx] = votes.get(idx, 0) + 1

        # Filter by voting threshold.
        filtered = [
            idx for idx, count in votes.items()
            if count >= voting_threshold
        ]

        if not filtered:
            # Fall back to all candidates if threshold too strict.
            filtered = list(votes.keys())

        # Compute exact distances for candidates.
        candidate_indices = np.array(filtered)
        distances = np.linalg.norm(
            self.data[candidate_indices] - q, axis=1
        )

        # Return top k.
        top_k = np.argsort(distances)[:k]
        return candidate_indices[top_k], distances[top_k]


# Example usage:
np.random.seed(42)
data = np.random.randn(10000, 128).astype(np.float32)
query = np.random.randn(128).astype(np.float32)

index = MRPTIndex(n_trees=20, max_leaf_size=30)
index.fit(data)

indices, dists = index.query(query, k=5, voting_threshold=2)
print(f"Nearest neighbors: {indices}")
print(f"Distances: {dists}")

Autotuning the Voting Threshold

This snippet demonstrates how to automatically select the voting threshold for a target recall using a sample of queries.

def autotune_threshold(
    index: MRPTIndex,
    data: np.ndarray,
    n_sample: int = 100,
    k: int = 10,
    target_recall: float = 0.9,
) -> int:
    """
    Find the highest voting threshold that achieves the target recall.
    Higher threshold = fewer candidates = faster queries.
    """
    sample_idx = np.random.choice(len(data), n_sample, replace=False)

    # Compute brute-force ground truth for samples.
    ground_truth = {}
    for idx in sample_idx:
        q = data[idx]
        dists = np.linalg.norm(data - q, axis=1)
        dists[idx] = np.inf  # exclude self
        ground_truth[idx] = set(np.argsort(dists)[:k])

    best_threshold = 1
    for v in range(index.n_trees, 0, -1):
        recalls = []
        for idx in sample_idx:
            result_idx, _ = index.query(data[idx], k=k, voting_threshold=v)
            result_set = set(result_idx) - {idx}
            recall = len(result_set & ground_truth[idx]) / k
            recalls.append(recall)

        avg_recall = np.mean(recalls)
        if avg_recall >= target_recall:
            best_threshold = v
            break

    return best_threshold


optimal_v = autotune_threshold(index, data, target_recall=0.9)
print(f"Optimal voting threshold for 90% recall: {optimal_v}")

Common Pitfalls

1. Using too few trees. A single random projection tree has poor recall because nearby points often end up in different leaves. The ensemble is what makes MRPT work. For most datasets, 10–50 trees are needed to achieve 90%+ recall. Monitor recall on a validation set rather than guessing the tree count.
2. Setting the voting threshold without validation. A threshold that’s too high eliminates true neighbors; one that’s too low returns too many candidates, negating the speed benefit. Always use autotuning or at minimum validate against brute-force results on a sample.
3. Ignoring the curse of dimensionality. For very high-dimensional data (d > 1000), random projections become less effective at discriminating between points—all pairwise distances converge. Consider dimensionality reduction (PCA) as a preprocessing step, or switch to methods like HNSW that are more robust in high dimensions.
4. Not normalizing vectors. If the dataset contains vectors of varying norms, the random projection splits will be dominated by the high-norm points. Normalize vectors to unit length before indexing if you are using cosine similarity as your distance metric.
5. Rebuilding the entire index for small updates. While MRPT supports incremental insertion, avoid rebuilding all trees when only a few points change. Use a tiered approach: accumulate new points in a small brute-force buffer and periodically merge into the main index with a full rebuild.

Key Takeaways