feat: product quantization fast scan

[usamoi]

blocked by #537

[VoVAllen]

Can you show us some performance numbers?

[usamoi]

dataset: glove-100-angular.hdf5 (1183514 vectors, 100 dims, 10000 tests, L2 normalized, IP distance)
machine: CPU: AMD Ryzen 7 6800H (16) @ 4.7GHz, Arch Linux x86_64, 6.10.2-arch1-1, 30.08 GiB Memory

Faiss here is https://pypi.org/project/faiss-cpu/.

Faiss-IVF (nlist = 1000; nprobe = 10): recall = 0.819589999999983, time = 11.370838165283203 s
Faiss-PQ (ratio = 2, bits = 4; fast_scan = off, rerank_size = 100): (timeout)
Faiss-PQ (ratio = 2, bits = 4; fast_scan = on, rerank_size = 100): recall = 0.9166199999999625, time = 13.678839206695557 s
Faiss-IVFPQ (nlist = 1000, ratio = 2, bits = 4; nprobe = 10, fast_scan = off, rerank_size = 100): recall = 0.803549999999983, time = 59.124324321746826 s
Faiss-IVFPQ (nlist = 1000, ratio = 2, bits = 4; nprobe = 10, fast_scan = on, rerank_size = 100): recall = 0.7459399999999857, time = 0.5915257930755615 s

pgvecto.rs-IVF (nlist = 1000; nprobe = 10): recall = 0.7628899999999892, time = 4.944989 s
pgvecto.rs-PQ (ratio = 2, bits = 4; fast_scan = off, rerank_size = 100): (timeout)
pgvecto.rs-PQ (ratio = 2, bits = 4; fast_scan = on, rerank_size = 100): recall = 0.9271499999999622, time = 118.084755 s
pgvecto.rs-IVFPQ (nlist = 1000, ratio = 2, bits = 4; nprobe = 10, fast_scan = off, rerank_size = 100): recall = 0.7379399999999887, time = 8.469601 s
pgvecto.rs-IVFPQ (nlist = 1000, ratio = 2, bits = 4; nprobe = 10, fast_scan = on, rerank_size = 100): recall = 0.7376499999999886, time = 2.937704 s

[usamoi]

@VoVAllen

The result shows that fast-scan speeds pgvecto.rs a lot (2x ~ 3x) but the effect is not as excellent as that of faiss (64x).

The flamegraph (for pgvecto.rs-PQ, ratio = 2, bits = 4; fast_scan = on, rerank_size = 100) shows quantization::fast_scan::bit4::fast_scan_v3 only takes up 8.35% time in a query.

• fast_scan takes up 8.35%
• make_heap takes up 46.66%; it does not exist in faiss but it's needed by VBASE
• memory allocation takes up 13.60%; it does not exist in faiss
• dequantizing rough distance of u16 type takes up 12.99%; it can be optimized, and I think it does not exist in faiss

[VoVAllen]

• I'm surprised that for IVF the bottleneck is at heap instead of computation. Can you try a dataset with larger vector dim like GIST?
• Why PQ on full vector takes 100s+? (pgvecto.rs-PQ (ratio = 2, bits = 4; fast_scan = on, rerank_size = 100): recall = 0.9271499999999622, time = 118.084755 s)

[usamoi]

• I'm surprised that for IVF the bottleneck is at heap instead of computation. Can you try a dataset with larger vector dim like GIST?

It's expected by me since computation is too fast if fast scan is enabled (float32 in look-up table is quantized to int8, and AVX2 processes 32 vectors in a time). The main difference between our implementation and faiss is whether heap is used.

• Why PQ on full vector takes 100s+?

What do you mean? Is it too fast or too slow?

[VoVAllen]

Faiss-PQ (ratio = 2, bits = 4; fast_scan = on, rerank_size = 100): recall = 0.9166199999999625, time = 13.678839206695557 s

It's seems much slower than faiss, but other numbers are faster than faiss. I don't understand the problem here.

[VoVAllen]

Would it be better if the heap is O(1) for insertion and O(logN) for pop operation?

Reference:

• https://github.com/mpdn/radix-heap
• https://github.com/hanmertens/dary_heap
• https://github.com/conradludgate/fibonacii-heap

[usamoi]

pgvecto.rs-IVFPQ (fast_scan = on) is much slower than Faiss-IVFPQ (fast_scan = on) too.

[VoVAllen]

Another way to improve the heap operation might be to do it in a batch, which won't pollute cache less at each insertion.

[usamoi]

Would it be better if the heap is O(1) for insertion and O(logN) for pop operation?

There is no heap insertion but only make_heap. make_heap is O(n).

I'll try adding fast-path for make_heap: if there are less than 100 pops, do not make_heap, but maintain a small 100-sized heap.

[VoVAllen]

100 dim vector is rare to see in Deep Learning world. Larger dims is more typical. For 1000dim, the cost of the heap might not account for a large proportion.

[usamoi]

I'm merging it now because it blocks #549