百行Julia代码实现高效向量搜索

感谢 HN 用户 mik1998 和 borodi 指出 popcnt 指令，以及 Julia 中执行该指令的 count_ones 函数。我已经更新了计时结果，现在速度甚至更快了。

无论使用 Int8、UInt8 还是 Int64，在执行实际搜索时，性能都是相同的（在我的 M1 Air 上）。具体结果可能会因所使用的架构而异。最佳表示形式很可能是 UInt64，因为它或许能让你在不同架构上获得最具可迁移性的结果。

最佳搜索耗时来自以下两种方案之一：1）由 StaticArrays 组成的 Vector，加上 StaticArray 查询向量；2）元素按 (num_bytes, num_rows) 排列的 Matrix，加上 StaticArray 查询向量。因此，只需将查询数组设为静态数组，也能获得相同的性能。

这是一次实验，旨在探索使用二进制向量空间能以多快的速度执行精确的 RAG 查找。

“检索增强生成”（Retrieval Augmented Generation，RAG）是对 LLM 内容生成的增强。其基本架构是使用提示词和查询调用 LLM API。

但有时，我们希望通过编程方式向提示词中添加上下文，以便更好地回答查询。当我们想要查询 LLM 本不应了解的内部文档，或者只是想提升 LLM 在特定领域中的问答能力时，这尤其有用。

RAG 是确定应该添加哪些上下文的过程。通常的做法是将查询编码到一个向量空间中，然后在数据库中查找与查询最接近的向量。每个向量代表数百或数千个 token。

为什么选择二进制向量空间？

事实证明，它的准确率与完整的 32 位向量非常接近。但它能大幅节省服务器成本，也让内存检索变得更加可行。

一个原本为 1TB 的数据库现在约为 32GB，可以轻松装入成本低得多的配置的 RAM 中。

假设每一行都是一个 1024 维的 float32 向量，那么每行占 4096 字节，也就是约 2.44 亿行。

再补充一些背景：Wikipedia 可以说是公开可用于 RAG 的最大规模数据，但可能也只有约 1500 万行。

https://en.wikipedia.org/wiki/Wikipedia:Size_of_Wikipedia

每篇文章平均 668 个单词（共 45.5 亿个单词）

根据 OpenAI 的说法，750 个单词大约相当于 1000 个 token。因此，我们按 60 亿个 token 计算。

如果使用 mixedbread 模型，可以用一个向量表示 512 个 token。假设存在重叠，那么可以认为每个向量包含 400 个新 token，另外 112 个 token 用于重叠。

60 亿 / 400 = 约 15,000,000 个向量

内部 RAG 的规模很可能会远小于这个数字。

于是，这引出了一个问题：

“数据集要大到什么程度，针对二进制向量进行精确的暴力搜索才会变得完全不可行？”

如果我们的数据集比这个规模更小，那么只需用很少的代码采用最朴素的方案，就完全没有问题。

我们可能并不需要复杂的向量数据库，也不需要最新的 HSNW 库（HNSW 本身就是近似搜索，而我们做的是精确搜索）。

我选择使用 Julia，是因为：

在这项任务上，它应该能够获得与 C、C++ 或 Rust 相当的性能。因此，我们不会错失潜在的性能收益。

CODE BLOCKCOPYjulia> versioninfo()
Julia Version 1.11.0-beta1
Commit 08e1fc0abb9 (2024-04-10 08:40 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (arm64-apple-darwin22.4.0)
  CPU: 8 × Apple M1
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, apple-m1)
Threads: 4 default, 0 interactive, 2 GC (on 4 virtual cores)
Environment:
  JULIA_STACKTRACE_MINIMAL = true
  DYLD_LIBRARY_PATH = /Users/lunaticd/.wasmedge/lib
  JULIA_EDITOR = nvim

julia> versioninfo()
Julia Version 1.11.0-beta1
Commit 08e1fc0abb9 (2024-04-10 08:40 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (arm64-apple-darwin22.4.0)
  CPU: 8 × Apple M1
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, apple-m1)
Threads: 4 default, 0 interactive, 2 GC (on 4 virtual cores)
Environment:
  JULIA_STACKTRACE_MINIMAL = true
  DYLD_LIBRARY_PATH = /Users/lunaticd/.wasmedge/lib
  JULIA_EDITOR = nvim

每个向量都由 n 个 bit 组成，其中 n 是嵌入的维度。这个值会随所使用的嵌入模型而变化，目前有多种专有和开源实现。这里我假设使用的是 mixedbread 模型：mixedbread-ai/mxbai-embed-large-v1

该模型返回一个包含 1024 个元素的 float32 向量。经过二值化后，它会被转换成一个 1024 bit 的向量，然后表示为一个包含 128 个元素的 int8 向量，也就是 128 字节。此外，它还可以进一步缩减到 64 字节，这也是我们最终使用的表示形式。

一个 512 bit 的向量，表示为 64 个 int8 或 uint8 元素。

CODE BLOCKCOPYjulia> x = rand(Int8, 64)
64-element Vector{Int8}:
   26
  123
   70
   -2
    ⋮

julia> x = rand(Int8, 64)
64-element Vector{Int8}:
   26
  123
   70
   -2
    ⋮

请记住，虽然它以字节为单位表示，但我们的操作是在 bit 层面进行的。

比较两个 bit 向量需要使用汉明距离。幸运的是，它的定义非常直观：向量中每出现一个不匹配的 bit，就将距离加 1。因此，两个长度为 512 且完全不同的 bit 向量，其汉明距离为 512；如果它们完全相同，则距离为 0。

一种简单的实现方式是将整数表示转换为 bit 字符串，然后计算二者之间的距离。

CODE BLOCKCOPYjulia> bitstring(Int8(125))
"01111101"

julia> bitstring(Int8(33))
"00100001"

julia> function hamming_distance(s1::AbstractString, s2::AbstractString)::Int
           s = 0
           for (c1, c2) in zip(s1, s2)
               if c1 != c2
                   s += 1
               end
           end
           s
       end

julia> hamming_distance(bitstring(Int8(125)), bitstring(Int8(33))) # should be 4
4

julia> bitstring(Int8(125))
"01111101"

julia> bitstring(Int8(33))
"00100001"

julia> function hamming_distance(s1::AbstractString, s2::AbstractString)::Int
           s = 0
           for (c1, c2) in zip(s1, s2)
               if c1 != c2
                   s += 1
               end
           end
           s
       end

julia> hamming_distance(bitstring(Int8(125)), bitstring(Int8(33))) # should be 4
4

Julia 有一个很棒的基准测试包，名为 Chairmarks。让我们看看这个实现有多快。

CODE BLOCKCOPY
julia> using Chairmarks

julia> @be hamming_distance(bitstring(Int8(125)), bitstring(Int8(33)))
Benchmark: 6053 samples with 317 evaluations
min    36.672 ns (4 allocs: 128 bytes)
median 38.514 ns (4 allocs: 128 bytes)
mean   46.525 ns (4 allocs: 128 bytes, 0.13% gc time)
max    8.737 μs (4 allocs: 128 bytes, 98.30% gc time)

julia> using Chairmarks

julia> @be hamming_distance(bitstring(Int8(125)), bitstring(Int8(33)))
Benchmark: 6053 samples with 317 evaluations
min    36.672 ns (4 allocs: 128 bytes)
median 38.514 ns (4 allocs: 128 bytes)
mean   46.525 ns (4 allocs: 128 bytes, 0.13% gc time)
max    8.737 μs (4 allocs: 128 bytes, 98.30% gc time)

说实话，这已经相当不错了。作为对比，我们来做一次简单的加法：

CODE BLOCKCOPYjulia> @be 125 + 33
Benchmark: 7274 samples with 8775 evaluations
min    1.239 ns
median 1.339 ns
mean   1.453 ns
max    6.990 ns

julia> @be 125 + 33
Benchmark: 7274 samples with 8775 evaluations
min    1.239 ns
median 1.339 ns
mean   1.453 ns
max    6.990 ns

对于一次操作而言，我们不可能得到比这更低的结果。在 CPU 领域，纳秒基本上就是时间的原子测量单位。

不过，这并不意味着每个操作都必须花费约 1ns。硬件可以进行多种优化，使每个时钟周期内完成多个操作。CPU 的每个核心都有若干算术单元，用于处理基本操作（加法、减法、xor、or、and 等），而算术单元通过所谓的流水线机制获得指令。关键在于，算术单元不止一个，因此只要这些操作彼此不依赖，理论上我们就可以在同一个时钟周期内执行多次加法或 bit 操作。让我们只使用基本操作重写 hamming_distance，这样或许就能利用这一点。

考虑比较两个 bit：

这就是 XOR 操作。我们希望对所有 bit 执行 XOR，记录其中哪些结果为 1，然后返回它们的总和。

为了对所有 bit 完成这一操作，可以按 bit 数对整数进行移位。对于一个 8-bit 数字，需要执行 7 次移位。

具体实现如下：

CODE BLOCKCOPY@inline function hamming_distance(x1::T, x2::T)::Int where {T<:Union{Int8,UInt8}}
    r = x1 ⊻ x2 # xor
    # how many xored bits are 1
    c = 0
    for i in 0:7
        c += (r >> i) & 1
    end
    return Int(c)
end

# benchmark again
julia> @be hamming_distance(Int8(33), Int8(125))
Benchmark: 4797 samples with 8966 evaluations
min    2.175 ns
median 2.189 ns
mean   2.205 ns
max    5.279 ns

@inline function hamming_distance(x1::T, x2::T)::Int where {T<:Union{Int8,UInt8}}
    r = x1 ⊻ x2 # xor
    # how many xored bits are 1
    c = 0
    for i in 0:7
        c += (r >> i) & 1
    end
    return Int(c)
end

# benchmark again
julia> @be hamming_distance(Int8(33), Int8(125))
Benchmark: 4797 samples with 8966 evaluations
min    2.175 ns
median 2.189 ns
mean   2.205 ns
max    5.279 ns

这是一个巨大的改进，几乎和执行一次加法运算一样快！

Julia 有一个内置函数 count_ones，它会被转换为机器指令 popcnt 来统计所有的 1，所以我们可以直接使用它，而不需要做移位操作！

@inline function hamming_distance(x1::T, x2::T)::Int where {T<:Union{Int8,UInt8}}
    return Int(count_ones(x1 ⊻ x2))
end

现在它的性能和执行一次加法运算一样！

julia> @be hamming_distance(Int8(33), Int8(125))
Benchmark: 5800 samples with 13341 evaluations
min    1.243 ns
median 1.249 ns
mean   1.251 ns
max    3.760 ns

现在我们需要为向量实现这个功能：

@inline function hamming_distance1(x1::AbstractArray{T}, x2::AbstractArray{T})::Int where {T<:Union{Int8,UInt8}}
    s = 0
    for i in 1:length(x1)
        s += hamming_distance(x1[i], x2[i])
    end
    s
end

我们预期这需要大约 128ns（~2 * 64）。

julia> @be hamming_distance1(q1, q2)
Benchmark: 4469 samples with 256 evaluations
min    69.500 ns
median 75.035 ns
mean   80.240 ns
max    189.941 ns

还有一些其他的方法可以进一步优化这个循环。编译器有时会自动做这些优化，但我们可以手动标注来确保一定会做。

@inline function hamming_distance(x1::AbstractArray{T}, x2::AbstractArray{T})::Int where {T<:Union{Int8,UInt8}}
    s = 0
    @inbounds @simd for i in eachindex(x1, x2)
        s += hamming_distance(x1[i], x2[i])
    end
    s
end

@inbounds 移除任何边界检查。

@simd 告诉编译器尽可能使用 SIMD 指令。

julia> @be hamming_distance(q1, q2)
Benchmark: 4441 samples with 323 evaluations
min    52.759 ns
median 56.245 ns
mean   61.721 ns
max    163.827 ns

这是一个不错的改进。这个基准在一个全新的 REPL 会话中可以达到 < 20ns。我认为这取决于当前的后台使用情况以及 SIMD 在那个时刻的利用效率。

julia> q1 = rand(Int8, 64);

julia> q2 = rand(Int8, 64);

julia> hamming_distance(q1, q2);

julia> @be hamming_distance1(q1, q2)
Benchmark: 4381 samples with 1113 evaluations
min    17.932 ns
median 18.381 ns
mean   19.253 ns
max    54.956 ns

在完美的环境中（比如一个隔离的机器上），你可以真的充分利用硬件。

假设每次比较需要 50ns，那么我们可以在一秒内搜索 2000 万行，或者如果使用 4 核的话可以搜索 8000 万行。

100 万行应该在 4 核上耗时 12ms。维基百科的 1500 万行需要 180ms。

在实践中，它比这个速度还要快：

mutable struct MaxHeap
    const data::Vector{Pair{Int,Int}}
    current_idx::Int # add pairs until current_idx > length(data)
    const k::Int

function MaxHeap(k::Int)
        new(fill((typemax(Int) => -1), k), 1, k)
    end
end

function insert!(heap::MaxHeap, value::Pair{Int,Int})
    if heap.current_idx <= heap.k
        heap.data[heap.current_idx] = value
        heap.current_idx += 1
        if heap.current_idx > heap.k
            makeheap!(heap)
        end
    elseif value.first < heap.data[1].first
        heap.data[1] = value
        heapify!(heap, 1)
    end
end

function makeheap!(heap::MaxHeap)
    for i in div(heap.k, 2):-1:1
        heapify!(heap, i)
    end
end

function heapify!(heap::MaxHeap, i::Int)
    left = 2 * i
    right = 2 * i + 1
    largest = i

if left <= length(heap.data) && heap.data[left].first > heap.data[largest].first
        largest = left
    end

if right <= length(heap.data) && heap.data[right].first > heap.data[largest].first
        largest = right
    end
end

堆的结构可以忽略，它只是一个最大堆，有一个最大大小限制。被搜索的数据库的每个部分都有自己的堆，然后在最后排序并选择前 k 个结果。

julia> X1 = [rand(Int8, 64) for _ in 1:(10^6)];

julia> k_closest(X1, q1, 1)
1-element Vector{Pair{Int64, Int64}}:
 202 => 76839

julia> @be k_closest(X1, q1, 1)
Benchmark: 17 samples with 1 evaluation
min    5.571 ms (3 allocs: 112 bytes)
median 5.618 ms (3 allocs: 112 bytes)
mean   5.853 ms (3 allocs: 112 bytes)
max    6.996 ms (3 allocs: 112 bytes)

julia> k_closest_parallel(X1, q1, 1)
1-element Vector{Pair{Int64, Int64}}:
 202 => 76839

julia> @be k_closest_parallel(X1, q1, 1)
Benchmark: 32 samples with 1 evaluation
min    3.022 ms (54 allocs: 3.672 KiB)
median 3.070 ms (54 allocs: 3.672 KiB)
mean   3.074 ms (54 allocs: 3.672 KiB)
max    3.153 ms (54 allocs: 3.672 KiB)

每次向量比较约 6ns 而不是 ~50ns。

但等等，还有更多！我们可以使用 StaticArrays 做得更好：

julia> using StaticArrays

julia> q1 = SVector{64, Int8}(rand(Int8, 64));

julia> X1 = [SVector{64, Int8}(rand(Int8, 64)) for _ in 1:(10^6)];

julia> @be k_closest(X1, q1, 1)
Benchmark: 22 samples with 1 evaluation
min    4.049 ms (3 allocs: 112 bytes)
median 4.245 ms (3 allocs: 112 bytes)
mean   4.453 ms (3 allocs: 112 bytes)
max    6.010 ms (3 allocs: 112 bytes)

julia> @be k_closest_parallel(X1, q1, 1)
Benchmark: 86 samples with 1 evaluation
min    1.120 ms (54 allocs: 3.672 KiB)
median 1.133 ms (54 allocs: 3.672 KiB)
mean   1.144 ms (54 allocs: 3.672 KiB)
max    1.311 ms (54 allocs: 3.672 KiB)

julia> @be k_closest(X1, q1, 1) Benchmark: 22 samples with 1 evaluation min 4.049 ms (3 allocs: 112 bytes) median 4.245 ms (3 allocs: 112 bytes) mean 4.453 ms (3 allocs: 112 bytes) max 6.010 ms (3 allocs: 112 bytes)

julia> @be k_closest_parallel(X1, q1, 1) Benchmark: 86 samples with 1 evaluation min 1.120 ms (54 allocs: 3.672 KiB) median 1.133 ms (54 allocs: 3.672 KiB) mean 1.144 ms (54 allocs: 3.672 KiB) max 1.311 ms (54 allocs: 3.672 KiB)

现在，每次向量比较仅耗时约 4ns！另外请注意，并行实现似乎更充分地利用了可用的核心，将运行时间改善了 4 倍，而不只是 2 倍，因为我们有 4 个核心。

不过，我们想要的不只是最接近的一个向量，而是最接近的“k”个向量，因此让我们把“k”设置为一个更符合实际情况的值：

CODE BLOCKCOPYjulia> @be k_closest(X1, q1, 20) Benchmark: 22 samples with 1 evaluation min 4.033 ms (5 allocs: 832 bytes) median 4.267 ms (5 allocs: 832 bytes) mean 4.511 ms (5 allocs: 832 bytes) max 5.910 ms (5 allocs: 832 bytes)

julia> @be k_closest_parallel(X1, q1, 20) Benchmark: 85 samples with 1 evaluation min 1.126 ms (56 allocs: 7.828 KiB) median 1.137 ms (56 allocs: 7.828 KiB) mean 1.142 ms (56 allocs: 7.828 KiB) max 1.194 ms (56 allocs: 7.828 KiB)

julia> @be k_closest(X1, q1, 100) Benchmark: 23 samples with 1 evaluation min 4.108 ms (5 allocs: 3.281 KiB) median 4.270 ms (5 allocs: 3.281 KiB) mean 4.344 ms (5 allocs: 3.281 KiB) max 5.472 ms (5 allocs: 3.281 KiB)

julia> @be k_closest_parallel(X1, q1, 100) Benchmark: 82 samples with 1 evaluation min 1.159 ms (58 allocs: 23.906 KiB) median 1.185 ms (58 allocs: 23.906 KiB) mean 1.201 ms (58 allocs: 23.906 KiB) max 1.384 ms (58 allocs: 23.906 KiB)

julia> @be k_closest(X1, q1, 20) Benchmark: 22 samples with 1 evaluation min 4.033 ms (5 allocs: 832 bytes) median 4.267 ms (5 allocs: 832 bytes) mean 4.511 ms (5 allocs: 832 bytes) max 5.910 ms (5 allocs: 832 bytes)

julia> @be k_closest_parallel(X1, q1, 20) Benchmark: 85 samples with 1 evaluation min 1.126 ms (56 allocs: 7.828 KiB) median 1.137 ms (56 allocs: 7.828 KiB) mean 1.142 ms (56 allocs: 7.828 KiB) max 1.194 ms (56 allocs: 7.828 KiB)

julia> @be k_closest(X1, q1, 100) Benchmark: 23 samples with 1 evaluation min 4.108 ms (5 allocs: 3.281 KiB) median 4.270 ms (5 allocs: 3.281 KiB) mean 4.344 ms (5 allocs: 3.281 KiB) max 5.472 ms (5 allocs: 3.281 KiB)

julia> @be k_closest_parallel(X1, q1, 100) Benchmark: 82 samples with 1 evaluation min 1.159 ms (58 allocs: 23.906 KiB) median 1.185 ms (58 allocs: 23.906 KiB) mean 1.201 ms (58 allocs: 23.906 KiB) max 1.384 ms (58 allocs: 23.906 KiB)

由于我们使用堆来对条目进行排序，因此耗时相近。

有趣的是，让我们在 1500 万行数据（Wikipedia）上搜索一下。

CODE BLOCKCOPYjulia> X1 = [SVector{64, Int8}(rand(Int8, 64)) for _ in 1:(15 * 10^6)];

julia> @be k_closest_parallel(X1, q1, 10) Benchmark: 6 samples with 1 evaluation min 16.534 ms (56 allocs: 5.625 KiB) median 16.637 ms (56 allocs: 5.625 KiB) mean 16.628 ms (56 allocs: 5.625 KiB) max 16.762 ms (56 allocs: 5.625 KiB)

julia> @be k_closest_parallel(X1, q1, 50) Benchmark: 6 samples with 1 evaluation min 16.506 ms (58 allocs: 14.062 KiB) median 16.672 ms (58 allocs: 14.062 KiB) mean 16.669 ms (58 allocs: 14.062 KiB) max 16.777 ms (58 allocs: 14.062 KiB)

julia> X1 = [SVector{64, Int8}(rand(Int8, 64)) for _ in 1:(15 * 10^6)];

julia> @be k_closest_parallel(X1, q1, 10) Benchmark: 6 samples with 1 evaluation min 16.534 ms (56 allocs: 5.625 KiB) median 16.637 ms (56 allocs: 5.625 KiB) mean 16.628 ms (56 allocs: 5.625 KiB) max 16.762 ms (56 allocs: 5.625 KiB)

julia> @be k_closest_parallel(X1, q1, 50) Benchmark: 6 samples with 1 evaluation min 16.506 ms (58 allocs: 14.062 KiB) median 16.672 ms (58 allocs: 14.062 KiB) mean 16.669 ms (58 allocs: 14.062 KiB) max 16.777 ms (58 allocs: 14.062 KiB)

再来看看 10 万行数据：

CODE BLOCKCOPYjulia> X1 = [SVector{64, Int8}(rand(Int8, 64)) for _ in 1:(10^5)];

julia> @be k_closest_parallel(X1, q1, 30) Benchmark: 679 samples with 1 evaluation min 120.250 μs (56 allocs: 10.000 KiB) median 131.792 μs (56 allocs: 10.000 KiB) mean 137.709 μs (56 allocs: 10.000 KiB) max 463.625 μs (56 allocs: 10.000 KiB)

julia> X1 = [SVector{64, Int8}(rand(Int8, 64)) for _ in 1:(10^5)];

julia> @be k_closest_parallel(X1, q1, 30) Benchmark: 679 samples with 1 evaluation min 120.250 μs (56 allocs: 10.000 KiB) median 131.792 μs (56 allocs: 10.000 KiB) mean 137.709 μs (56 allocs: 10.000 KiB) max 463.625 μs (56 allocs: 10.000 KiB)

usearch 似乎是内存向量相似度搜索领域当前最先进的实现。

CODE BLOCKCOPYIn [8]: from usearch.index import search, MetricKind ...: import numpy as np

In [12]: X = np.random.randint(-128, 128, size=(10**6, 64), dtype=np.int8)

In [13]: q = np.random.randint(-128, 128, 64, dtype=np.int8)

In [14]: %timeit search(X, q, 1, MetricKind.Hamming, exact=True) 62.4 ms ± 494 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [15]: %timeit search(X, q, 1, MetricKind.Hamming, exact=True, threads=4) 32.7 ms ± 1.7 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [8]: from usearch.index import search, MetricKind ...: import numpy as np

In [12]: X = np.random.randint(-128, 128, size=(10**6, 64), dtype=np.int8)

In [13]: q = np.random.randint(-128, 128, 64, dtype=np.int8)

In [14]: %timeit search(X, q, 1, MetricKind.Hamming, exact=True) 62.4 ms ± 494 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [15]: %timeit search(X, q, 1, MetricKind.Hamming, exact=True, threads=4) 32.7 ms ± 1.7 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Julia 实现在单线程下快约 16 倍，使用 4 个核心时快约 30 倍。

也许还有办法让 usearch 变得更快；公平地说，尽可能提升精确搜索的速度很可能并不是它的优先目标。

话虽如此，这只有大约 100 行 Julia 代码，而且除了使用 StaticArrays——这一点或许也算——之外，我们并没有做任何花哨的事情。

仅用约 100 行 Julia 代码，我们就能够达到最先进的水平