Asymmetric LSH in Python

Uses different transformations for data and queries to reduce maximum inner product search to LSH for angles.

File: asymmetric_lsh.py, 90 lines
Level: I0 not checked yet
Toolchain: Python 3.14, pytest 9
Test framework: pytest
Run the tests: python3 -m pytest

asymmetric_lsh.pyRaw GitHub

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/asymmetric-lsh/. The authors are listed in prov/NNNB37.jsonl.
"""Asymmetric LSH for maximum inner product search (the simple transformation with random hyperplanes)."""

import math
import random
from typing import Optional, Sequence

Vector = Sequence[float]


def dot(a: Vector, b: Vector) -> float:
    return sum(x * y for x, y in zip(a, b))


def norm(a: Vector) -> float:
    return math.sqrt(dot(a, a))


def transform_data(x: Vector, m: float) -> list[float]:
    """P(x) = [x/M ; sqrt(1 - |x|^2/M^2)], a unit vector when |x| <= M."""
    y = [v / m for v in x]
    return y + [math.sqrt(max(0.0, 1.0 - dot(y, y)))]


def transform_query(q: Vector) -> list[float]:
    """Q(q) = [q/|q| ; 0], a unit vector."""
    s = norm(q)
    return [v / s for v in q] + [0.0]


def l2_alsh_transforms(x: Vector, q: Vector, m: int) -> tuple[list[float], list[float]]:
    """The L2-ALSH maps: P(x) = [x; |x|^2; |x|^4; ...; |x|^(2^m)], Q(q) = [q; 1/2; ...; 1/2].

    x must already be scaled so that |x| < 1, and q normalized to |q| = 1.
    """
    s = dot(x, x)
    tail = []
    for _ in range(m):
        tail.append(s)
        s = s * s
    return list(x) + tail, list(q) + [0.5] * m


class SimpleALSH:
    def __init__(self, points: Sequence[Vector], k: int, L: int,
                 rng: Optional[random.Random] = None) -> None:
        rng = rng or random.Random()
        self.points = [list(p) for p in points]
        self.k, self.L = k, L
        self.m = max((norm(p) for p in self.points), default=0.0) or 1.0  # @line L1
        dim = len(self.points[0]) + 1 if self.points else 1
        self.planes = [[[rng.gauss(0.0, 1.0) for _ in range(dim)] for _ in range(k)]
                       for _ in range(L)]  # @line L3
        self.tables: list[dict[tuple[bool, ...], list[int]]] = [{} for _ in range(L)]
        for i, p in enumerate(self.points):
            px = transform_data(p, self.m)  # @line L2
            for j in range(L):  # @line L4
                self.tables[j].setdefault(self.key(j, px), []).append(i)

    def key(self, j: int, v: Vector) -> tuple[bool, ...]:
        return tuple(dot(a, v) >= 0.0 for a in self.planes[j])

    def candidates(self, q: Vector) -> set[int]:
        qq = transform_query(q)  # @line L5
        cand: set[int] = set()
        for j in range(self.L):  # @line L6
            cand.update(self.tables[j].get(self.key(j, qq), ()))
        return cand

    def query(self, q: Vector) -> Optional[int]:
        """Index of the candidate with the largest inner product with q, or None."""
        cand = self.candidates(q)
        if not cand:
            return None
        return max(sorted(cand), key=lambda i: dot(q, self.points[i]))  # @line L7


def collision_probability(s: float) -> float:
    """Pr[one hyperplane bit agrees] for unit vectors with inner product s."""
    return 1.0 - math.acos(max(-1.0, min(1.0, s))) / math.pi


def parameters(n: int, s: float, c: float) -> tuple[int, int, float]:
    """k, L and rho for an (S, cS) guarantee on the normalized inner product, success >= 1 - 1/e."""
    p1, p2 = collision_probability(s), collision_probability(c * s)
    rho = math.log(1 / p1) / math.log(1 / p2)
    k = max(1, math.ceil(math.log(n) / math.log(1 / p2)))
    L = math.ceil(n ** rho / p1)
    return k, L, rho

test_asymmetric_lsh.py 98 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/asymmetric-lsh/. The authors are listed in prov/NNNB37.jsonl.

import math
import random

from asymmetric_lsh import (SimpleALSH, collision_probability, dot, l2_alsh_transforms,
                            norm, parameters, transform_data, transform_query)

X = [(2, 0), (1, 1), (0, 3), (-1, -1)]


def test_example() -> None:
    q = (1, 1)
    m = max(norm(x) for x in X)
    assert m == 3
    sims = [dot(transform_data(x, m), transform_query(q)) for x in X]
    assert [round(s, 4) for s in sims] == [0.4714, 0.4714, 0.7071, -0.4714]
    assert round(collision_probability(sims[2]), 4) == 0.75
    assert round(collision_probability(sims[0]), 4) == 0.6563
    assert parameters(10 ** 6, 0.9, 0.5) == (32, 165, parameters(10 ** 6, 0.9, 0.5)[2])
    assert abs(parameters(10 ** 6, 0.9, 0.5)[2] - 0.3579) < 1e-4
    # One table with zero hyperplanes puts everything in one bucket: exact MIPS.
    assert SimpleALSH(X, k=0, L=1).query(q) == 2
    assert SimpleALSH([], k=2, L=2).query(q) is None


def test_transforms_random() -> None:
    rng = random.Random(37)
    for _ in range(500):
        d = rng.randint(1, 6)
        pts = [[rng.uniform(-5, 5) for _ in range(d)] for _ in range(rng.randint(1, 8))]
        q = [rng.uniform(-5, 5) for _ in range(d)]
        if norm(q) < 1e-6:
            continue
        m = max(norm(x) for x in pts)
        for x in pts:
            px, qq = transform_data(x, m), transform_query(q)
            assert abs(norm(px) - 1) < 1e-9 and abs(norm(qq) - 1) < 1e-9
            assert abs(dot(px, qq) - dot(q, x) / (m * norm(q))) < 1e-9
        # L2-ALSH identity: |Q - P|^2 = 1 + m/4 - 2 q.x + |x|^(2^(m+1)).
        u = rng.uniform(0.3, 0.95)
        x = [v * u / m for v in pts[0]]
        qn = [v / norm(q) for v in q]
        mm = rng.randint(1, 4)
        P, Q = l2_alsh_transforms(x, qn, mm)
        lhs = sum((a - b) ** 2 for a, b in zip(P, Q))
        rhs = 1 + mm / 4 - 2 * dot(qn, x) + dot(x, x) ** (2 ** mm)
        assert abs(lhs - rhs) < 1e-9


def test_query_against_oracle() -> None:
    rng = random.Random(3)
    for _ in range(500):
        d = rng.randint(1, 4)
        pts = [[rng.uniform(-3, 3) for _ in range(d)] for _ in range(rng.randint(1, 10))]
        idx = SimpleALSH(pts, k=rng.randint(0, 4), L=rng.randint(1, 4), rng=rng)
        q = [rng.uniform(-3, 3) for _ in range(d)]
        if norm(q) < 1e-6:
            continue
        cand = idx.candidates(q)
        got = idx.query(q)
        if not cand:
            assert got is None
            continue
        best = max(dot(q, pts[i]) for i in cand)
        assert got in cand and dot(q, pts[got]) == best
        # Candidates are exactly the points sharing a bucket in some table.
        qq = transform_query(q)
        m = max(norm(p) for p in pts)
        want = {i for i, p in enumerate(pts)
                if any(idx.key(j, transform_data(p, m)) == idx.key(j, qq) for j in range(idx.L))}
        assert cand == want


def test_collision_rate() -> None:
    rng = random.Random(11)
    u = transform_data([1.0, 2.0], 3.0)
    v = transform_query([2.0, 1.0])
    s = dot(u, v)
    trials = 4000
    hits = 0
    for _ in range(trials):
        a = [rng.gauss(0, 1) for _ in range(3)]
        hits += (dot(a, u) >= 0) == (dot(a, v) >= 0)
    assert abs(hits / trials - collision_probability(s)) < 0.03


def test_success_probability() -> None:
    rng = random.Random(5)
    n, ok, runs = 60, 0, 60
    for _ in range(runs):
        pts = [[rng.gauss(0, 1) for _ in range(4)] for _ in range(n)]
        q = [rng.gauss(0, 1) for _ in range(4)]
        best = max(range(n), key=lambda i: dot(q, pts[i]))
        idx = SimpleALSH(pts, k=3, L=12, rng=rng)
        ok += idx.query(q) == best
    assert ok >= runs * 0.8

What the next level needs

the tests did not run in this build

Read the full page about Asymmetric LSH