Locality-Sensitive Hashing (LSH)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296

"""
Random Hyperplane LSH for Cosine Similarity.
 
This implementation is optimized for efficiency while remaining readable.
Key optimizations:
1. Matrix operations instead of loops
2. Pre-computed random hyperplanes
3. Efficient bit packing for hash keys
"""
 
import numpy as np
from typing import List, Tuple, Optional, Dict, Set
from collections import defaultdict
from dataclasses import dataclass, field
import time
 
@dataclass
class RandomHyperplaneLSH:
    """
    LSH index using random hyperplanes for cosine similarity.
    
    Parameters:
    -----------
    dim : int
        Dimensionality of vectors
    num_hyperplanes : int (k)
        Number of hyperplanes per hash table. Higher k means:
        - Fewer candidates per bucket (more precise)
        - Higher risk of missing true neighbors
    num_tables : int (L)
        Number of hash tables. Higher L means:
        - More likely to find true neighbors
        - More candidates to check
        - More memory and query time
    seed : int
        Random seed for reproducibility
    """
    dim: int
    num_hyperplanes: int = 10  # k
    num_tables: int = 20      # L
    seed: int = 42
    
    # Initialized in __post_init__
    hyperplanes: np.ndarray = field(init=False)  # Shape: (L, k, dim)
    tables: List[Dict] = field(init=False)
    vectors: Optional[np.ndarray] = field(default=None, init=False)
    
    def __post_init__(self):
        np.random.seed(self.seed)
        
        # Generate random hyperplanes for all tables at once
        # Each hyperplane is a random normal vector
        self.hyperplanes = np.random.randn(
            self.num_tables, 
            self.num_hyperplanes, 
            self.dim
        )
        
        # Normalize hyperplanes (optional, but good practice)
        norms = np.linalg.norm(self.hyperplanes, axis=2, keepdims=True)
        self.hyperplanes = self.hyperplanes / norms
        
        # Initialize empty hash tables
        self.tables = [defaultdict(list) for _ in range(self.num_tables)]
    
    def _hash_vector(
        self, 
        vector: np.ndarray, 
        table_idx: int
    ) -> Tuple[int, ...]:
        """
        Compute hash key for a vector in a specific table.
        
        The hash is a tuple of 0s and 1s indicating which side
        of each hyperplane the vector lies on.
        """
        # Dot product with all hyperplanes in this table
        # Shape: (num_hyperplanes,)
        projections = np.dot(self.hyperplanes[table_idx], vector)
        
        # Convert to binary: 1 if positive, 0 if negative
        bits = (projections >= 0).astype(int)
        
        # Return as tuple (hashable)
        return tuple(bits.tolist())
    
    def _hash_batch(
        self, 
        vectors: np.ndarray
    ) -> List[List[Tuple[int, ...]]]:
        """
        Compute hash keys for multiple vectors across all tables.
        
        Returns: List of L lists, each containing n hash keys
        """
        n = len(vectors)
        all_hashes = []
        
        for table_idx in range(self.num_tables):
            # Compute all projections at once
            # Shape: (n, num_hyperplanes)
            projections = np.dot(vectors, self.hyperplanes[table_idx].T)
            
            # Convert to binary
            bits = (projections >= 0).astype(int)
            
            # Convert each row to tuple
            hashes = [tuple(row.tolist()) for row in bits]
            all_hashes.append(hashes)
        
        return all_hashes
    
    def fit(self, vectors: np.ndarray) -> 'RandomHyperplaneLSH':
        """
        Build the LSH index from a dataset of vectors.
        
        Parameters:
        -----------
        vectors : np.ndarray of shape (n, dim)
            Dataset of n vectors
        
        Returns:
        --------
        self : for method chaining
        """
        self.vectors = vectors.copy()
        n = len(vectors)
        
        # Normalize vectors (cosine similarity is scale-invariant)
        norms = np.linalg.norm(vectors, axis=1, keepdims=True)
        norms = np.where(norms == 0, 1, norms)  # Avoid division by zero
        normalized = vectors / norms
        
        # Compute all hashes
        all_hashes = self._hash_batch(normalized)
        
        # Insert into tables
        for table_idx in range(self.num_tables):
            for vec_idx in range(n):
                hash_key = all_hashes[table_idx][vec_idx]
                self.tables[table_idx][hash_key].append(vec_idx)
        
        return self
    
    def query(
        self, 
        query: np.ndarray, 
        num_neighbors: int = 5,
        return_similarity: bool = True
    ) -> List[Tuple[int, float]]:
        """
        Find approximate nearest neighbors for a query vector.
        
        Parameters:
        -----------
        query : np.ndarray of shape (dim,)
            Query vector
        num_neighbors : int
            Number of neighbors to return
        return_similarity : bool
            If True, return cosine similarity; else return angular distance
            
        Returns:
        --------
        List of (index, similarity/distance) tuples, sorted by similarity
        """
        if self.vectors is None:
            raise ValueError("Index not built. Call fit() first.")
        
        # Normalize query
        query_norm = np.linalg.norm(query)
        if query_norm == 0:
            raise ValueError("Query vector has zero norm")
        query_normalized = query / query_norm
        
        # Collect candidates from all tables
        candidates: Set[int] = set()
        for table_idx in range(self.num_tables):
            hash_key = self._hash_vector(query_normalized, table_idx)
            bucket = self.tables[table_idx].get(hash_key, [])
            candidates.update(bucket)
        
        if not candidates:
            # No candidates found; fall back to checking a sample
            candidates = set(
                np.random.choice(len(self.vectors), 
                                min(100, len(self.vectors)), 
                                replace=False)
            )
        
        # Compute exact cosine similarities for candidates
        candidate_indices = list(candidates)
        candidate_vectors = self.vectors[candidate_indices]
        
        # Normalize candidate vectors
        norms = np.linalg.norm(candidate_vectors, axis=1, keepdims=True)
        norms = np.where(norms == 0, 1, norms)
        candidate_normalized = candidate_vectors / norms
        
        # Compute similarities
        similarities = np.dot(candidate_normalized, query_normalized)
        
        # Sort by similarity (descending)
        sorted_indices = np.argsort(-similarities)
        
        # Return top-k
        results = []
        for i in sorted_indices[:num_neighbors]:
            idx = candidate_indices[i]
            sim = similarities[i]
            if return_similarity:
                results.append((idx, float(sim)))
            else:
                # Angular distance
                results.append((idx, float(np.arccos(np.clip(sim, -1, 1)) / np.pi)))
        
        return results
    
    def get_stats(self) -> Dict:
        """Get statistics about the index."""
        bucket_sizes = []
        for table in self.tables:
            bucket_sizes.extend(len(v) for v in table.values())
        
        return {
            "num_tables": self.num_tables,
            "num_hyperplanes": self.num_hyperplanes,
            "num_vectors": len(self.vectors) if self.vectors is not None else 0,
            "total_buckets": sum(len(t) for t in self.tables),
            "avg_bucket_size": np.mean(bucket_sizes) if bucket_sizes else 0,
            "max_bucket_size": max(bucket_sizes) if bucket_sizes else 0,
            "empty_buckets": sum(1 for t in self.tables for v in t.values() if len(v) == 0),
        }
 
 
def demo():
    """Demonstrate Random Hyperplane LSH."""
    np.random.seed(42)
    
    # Create dataset
    n = 50000
    dim = 256
    print(f"Creating dataset: {n} vectors, {dim} dimensions")
    
    vectors = np.random.randn(n, dim)
    
    # Create query (similar to vector 0)
    query = vectors[0] + 0.1 * np.random.randn(dim)
    
    # Build index
    print("\nBuilding LSH index...")
    start = time.time()
    lsh = RandomHyperplaneLSH(dim=dim, num_hyperplanes=12, num_tables=25)
    lsh.fit(vectors)
    build_time = time.time() - start
    print(f"Build time: {build_time:.3f}s")
    print(f"Index stats: {lsh.get_stats()}")
    
    # Query
    print("\nQuerying...")
    start = time.time()
    neighbors = lsh.query(query, num_neighbors=10)
    query_time = time.time() - start
    print(f"Query time: {query_time:.6f}s")
    
    print("\nTop 10 neighbors (LSH):")
    for idx, sim in neighbors:
        print(f"  Vector {idx}: similarity = {sim:.4f}")
    
    # Brute force comparison
    print("\nBrute force comparison...")
    start = time.time()
    query_norm = query / np.linalg.norm(query)
    norms = np.linalg.norm(vectors, axis=1, keepdims=True)
    normalized = vectors / norms
    all_sims = np.dot(normalized, query_norm)
    bf_time = time.time() - start
    
    top_10_bf = np.argsort(-all_sims)[:10]
    print(f"Brute force time: {bf_time:.6f}s")
    print("\nTop 10 neighbors (brute force):")
    for idx in top_10_bf:
        print(f"  Vector {idx}: similarity = {all_sims[idx]:.4f}")
    
    # Calculate speedup
    print(f"\nSpeedup: {bf_time / query_time:.1f}x")
    
    # Check recall
    lsh_set = {idx for idx, _ in neighbors}
    bf_set = set(top_10_bf)
    recall = len(lsh_set & bf_set) / len(bf_set)
    print(f"Recall@10: {recall:.1%}")
 
 
if __name__ == "__main__":
    demo()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

"""
Analysis of AND-OR amplification for random hyperplane LSH.
"""
 
import numpy as np
import matplotlib.pyplot as plt
 
def collision_prob(cos_sim: float) -> float:
    """
    Collision probability for a single hyperplane.
    
    p = 1 - θ/π where θ = arccos(cos_sim)
    """
    theta = np.arccos(np.clip(cos_sim, -1, 1))
    return 1 - theta / np.pi
 
def amplified_prob(p: float, k: int, L: int) -> float:
    """
    Probability after AND-OR amplification.
    
    AND: require all k hyperplanes to agree -> p^k
    OR: succeed if any of L tables match -> 1 - (1-p^k)^L
    """
    p_and = p ** k
    p_or = 1 - (1 - p_and) ** L
    return p_or
 
def analyze_amplification():
    """Analyze how k and L affect discrimination."""
    
    # Various similarity levels
    similarities = [0.95, 0.9, 0.8, 0.7, 0.5, 0.3, 0.1, 0.0]
    base_probs = [collision_prob(s) for s in similarities]
    
    print("Base collision probabilities:")
    print("-" * 50)
    for s, p in zip(similarities, base_probs):
        print(f"  cos_sim = {s:.2f}: p = {p:.4f}")
    
    # Example: Distinguish 0.9 similarity from 0.5 similarity
    p_near = collision_prob(0.9)   # ~0.856
    p_far = collision_prob(0.5)    # ~0.667
    
    print(f"\n\nTarget: Find vectors with sim >= 0.9, reject sim <= 0.5")
    print(f"Base probabilities: p_near = {p_near:.4f}, p_far = {p_far:.4f}")
    print(f"Base ratio: {p_near/p_far:.2f}")
    
    # Try different k and L combinations
    print("\n\nAmplification analysis (fixed memory budget):")
    print("-" * 70)
    print(f"{'k':>4} {'L':>4} {'P_near':>10} {'P_far':>10} {'Ratio':>10} {'FPs per 1000':>12}")
    print("-" * 70)
    
    for k in [5, 8, 10, 12, 15]:
        for L in [10, 20, 30, 50]:
            P_near = amplified_prob(p_near, k, L)
            P_far = amplified_prob(p_far, k, L)
            ratio = P_near / P_far if P_far > 0 else float('inf')
            fps = P_far * 1000  # Expected false positives per 1000 far points
            
            print(f"{k:>4} {L:>4} {P_near:>10.4f} {P_far:>10.4f} {ratio:>10.2f} {fps:>12.1f}")
    
    # Optimal for specific recall target
    print("\n\nFinding k, L for 99% recall with minimal false positives:")
    print("-" * 70)
    
    best_config = None
    best_fp_rate = float('inf')
    
    for k in range(4, 20):
        for L in range(5, 100):
            P_near = amplified_prob(p_near, k, L)
            P_far = amplified_prob(p_far, k, L)
            
            if P_near >= 0.99:  # 99% recall
                if P_far < best_fp_rate:
                    best_fp_rate = P_far
                    best_config = (k, L, P_near, P_far)
    
    if best_config:
        k, L, P_near, P_far = best_config
        print(f"Best config for 99% recall:")
        print(f"  k = {k}, L = {L}")
        print(f"  Recall (P_near) = {P_near:.4f}")
        print(f"  False positive rate (P_far) = {P_far:.4f}")
        print(f"  Space: {L} tables × {k} hyperplanes = {L*k} total hyperplanes")
 
 
if __name__ == "__main__":
    analyze_amplification()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

"""
Multi-probe extension for Random Hyperplane LSH.
"""
 
import numpy as np
from typing import Tuple, Set, List
from itertools import combinations
 
def generate_probes(
    base_hash: Tuple[int, ...], 
    num_flips: int = 2
) -> List[Tuple[int, ...]]:
    """
    Generate all hash keys within Hamming distance of base_hash.
    
    Parameters:
    -----------
    base_hash : tuple of ints
        The original hash key
    num_flips : int
        Maximum number of bits to flip
        
    Returns:
    --------
    List of hash keys to probe
    """
    probes = [base_hash]
    k = len(base_hash)
    
    # Generate all combinations of positions to flip
    for flips in range(1, num_flips + 1):
        for positions in combinations(range(k), flips):
            # Flip bits at these positions
            new_hash = list(base_hash)
            for pos in positions:
                new_hash[pos] = 1 - new_hash[pos]
            probes.append(tuple(new_hash))
    
    return probes
 
def multiprobe_query(
    query: np.ndarray,
    hyperplanes: np.ndarray,  # Shape: (num_tables, k, dim)
    tables: List[dict],
    num_probes: int = 5  # Number of probes per table
) -> Set[int]:
    """
    Multi-probe query: check nearby buckets in addition to exact match.
    
    Returns set of candidate indices.
    """
    candidates = set()
    num_tables = len(tables)
    
    for table_idx in range(num_tables):
        # Compute exact hash
        projections = np.dot(hyperplanes[table_idx], query)
        base_hash = tuple((projections >= 0).astype(int).tolist())
        
        # Generate probes
        # For efficiency, we generate probes based on projection magnitudes
        # Bits with small |projection| are more likely to flip
        abs_proj = np.abs(projections)
        flip_priority = np.argsort(abs_proj)  # Smallest first
        
        probes_checked = 0
        for flips in range(len(base_hash) + 1):
            if probes_checked >= num_probes:
                break
            
            for positions in combinations(flip_priority[:flips+3], flips):
                if probes_checked >= num_probes:
                    break
                    
                # Create probed hash
                probed_hash = list(base_hash)
                for pos in positions:
                    probed_hash[pos] = 1 - probed_hash[pos]
                probed_hash = tuple(probed_hash)
                
                # Get candidates from this bucket
                bucket = tables[table_idx].get(probed_hash, [])
                candidates.update(bucket)
                probes_checked += 1
    
    return candidates
 
# Example comparison
def compare_probe_strategies():
    """Compare single-probe vs multi-probe LSH."""
    np.random.seed(42)
    
    # Setup
    n, dim, k = 10000, 128, 12
    
    # Generate random hyperplanes
    hyperplanes = np.random.randn(1, k, dim)  # Single table
    hyperplanes /= np.linalg.norm(hyperplanes, axis=2, keepdims=True)
    
    # Generate data
    data = np.random.randn(n, dim)
    data /= np.linalg.norm(data, axis=1, keepdims=True)
    
    # Build table
    table = {}
    for i, vec in enumerate(data):
        proj = np.dot(hyperplanes[0], vec)
        h = tuple((proj >= 0).astype(int).tolist())
        if h not in table:
            table[h] = []
        table[h].append(i)
    
    # Query (find nearest to a specific vector)
    query = data[0] + 0.2 * np.random.randn(dim)
    query /= np.linalg.norm(query)
    
    # True 100-NN
    sims = np.dot(data, query)
    true_100nn = set(np.argsort(-sims)[:100])
    
    # Single-probe
    proj = np.dot(hyperplanes[0], query)
    exact_hash = tuple((proj >= 0).astype(int).tolist())
    single_candidates = set(table.get(exact_hash, []))
    single_recall = len(single_candidates & true_100nn) / 100
    
    # Multi-probe (5 probes)
    multi_candidates = multiprobe_query(query, hyperplanes, [table], num_probes=5)
    multi_recall = len(multi_candidates & true_100nn) / 100
    
    # Multi-probe (20 probes)
    multi_candidates_20 = multiprobe_query(query, hyperplanes, [table], num_probes=20)
    multi_recall_20 = len(multi_candidates_20 & true_100nn) / 100
    
    print("Recall comparison (single table, k=12):")
    print(f"  Single-probe: {single_recall:.1%} ({len(single_candidates)} candidates)")
    print(f"  Multi-probe (5): {multi_recall:.1%} ({len(multi_candidates)} candidates)")
    print(f"  Multi-probe (20): {multi_recall_20:.1%} ({len(multi_candidates_20)} candidates)")
 
 
if __name__ == "__main__":
    compare_probe_strategies()