Explore more granular vector quantization?

[kaivalnp]

Description

Today, Lucene supports 8, 4, 2 and 1 bit quantization.

Each quantization level typically has an upper bound of recall with exact KNN that it can produce (which is "exact KNN with quantized scores" v/s "exact KNN with original scores", see mikemccand/luceneutil#528) -- this is the information loss due to quantization itself (before approximate-ness from search algorithms like HNSW comes into picture).

Any algorithm operating on quantized scores alone cannot go beyond this recall (e.g. tweaking parameters like maxConn, beamWidth, fanout, etc. for HNSW) without using the original scores from un-quantized vectors for re-ranking -- which may not be feasible for some use cases (e.g. keeping the index in-memory for performance, where using un-quantized vectors increases memory footprint by ~4x in case of byte-quantized vectors).

In such cases, I wonder if Lucene could support more granular quantization options (say the equivalent of! 6-bit quantization) -- for more granular recall v/s memory requirements?

[kaivalnp]

However, I understand that packing vectors with arbitrary quantization levels + supporting various comparisons in Lucene is a challenge -- so I tried a slightly different approach:

What if vectors were first expanded to a higher dimension, and then quantized to one of the supported options? (e.g. represent 1024 dimension vectors in 1536 dimensions, then quantize to 4 bits: we get the memory-equivalent of 6-bit quantization)

A dimension "expansion" from X -> Y could be as simple as padding the original vector with (Y - X) zeros, and rotating the new Y dimension vector by an arbitrary rotation. We need a rotation because quantization occurs per-dimension, so any dimension with all zeros is not very useful -- we basically need to spread the information in X dimensions across Y, which is where a random rotation helps.

I took some 1024 dimension Cohere v3 docs + queries and "expanded" them to 1536 dimensions using the above strategy, which was a simple Python script like:

import numpy as np

def random_rotation_matrix(n):
    """Generate a random rotation matrix in n dimensions using QR decomposition."""
    H = np.random.randn(n, n)
    Q, R = np.linalg.qr(H)
    Q = Q @ np.diag(np.sign(np.diag(R)))
    return Q.astype("<f4")

def expand(size):
    rotation_matrix = random_rotation_matrix(size)
    def _func(arr):
        return np.matmul(np.pad(arr, (0, size - arr.size)), rotation_matrix)
    return _func

dim = 1024
num_docs = 100_000
input_vectors = "../cohere-v3-wikipedia-en-scattered-1024d.docs.first1M.vec"
docs = np.fromfile(doc_path, dtype=f"<{dim}f4", count=num_docs)

expanded_dim = 1536
expand_func = expand(expanded_dim)

docs_expanded = np.apply_along_axis(expand_func, 1, d)
docs_expanded.astype("<f4").tofile(f"cohere-v3-docs-{expanded_dim}d.vec")

# use the same expand_func for queries -- because it needs the same rotation as docs

Then I performed a recall test on the Cohere v3 vectors, numDocs = 100K, numQueries = 10K, forceMerge to 1 segment

Baseline (supported quantization options):

recall  latency(ms)  netCPU  avgCpuCount  quantized  visited  index(s)  index_docs/s  force_merge(s)  index_size(MB)  vec_disk(MB)  vec_RAM(MB)
 0.993        0.141   4.448       31.504         no     6858     19.29       5185.11           20.32          403.00       390.625      390.625
 0.983        0.088   2.760       31.288     8 bits     6868     20.35       4913.28           28.67          502.19       489.807       99.182
 0.927        0.059   1.822       31.037     4 bits     6954     20.59       4857.67           21.85          453.41       440.979       50.354
 0.835        0.037   1.152       30.970     2 bits     7299     20.23       4943.15           20.04          429.30       416.374       25.749
 0.717        0.032   0.976       30.599     1 bits     8197     19.86       5034.23           17.10          418.30       404.167       13.542

Candidate (non-standard quantization equivalent):

recall  latency(ms)  netCPU  avgCpuCount  quantized  visited  index(s)  index_docs/s  force_merge(s)  index_size(MB)  vec_disk(MB)  vec_RAM(MB)
 0.990        0.102   3.218       31.429    10 bits     6857     22.57       4429.88           40.30          624.26       611.877      123.596
 0.953        0.073   2.260       31.051     6 bits     6868     21.30       4694.62           32.77          673.11       660.706       74.768
 0.950        0.066   2.027       30.901     5 bits     6878     20.84       4798.69           29.12          563.26       550.842       62.561
 0.882        0.047   1.453       31.178     3 bits     7044     22.29       4486.12           28.81          636.63       623.894       37.956

The candidate uses "expanded" vectors of different dimensions + quantization options:

1. 10 bits == expand 1024 -> 1280 dimensions + 8 bit quantization
2. 6 bits == expand 1024 -> 1536 dimensions + 4 bit quantization
3. 5 bits == expand 1024 -> 1280 dimensions + 4 bit quantization
4. 3 bits == expand 1024 -> 1536 dimensions + 2 bit quantization

The exact KNN baseline in all runs is the one produced from the original 1024 dimension vectors.

It was encouraging to see the equivalent quantization options have about the correct interpolated values for recall and latency!

[kaivalnp]

This "expansion" to a higher dimension is probably better to do outside Lucene (ideally the model itself?), but sharing because I found it interesting that it works on existing vectors too!

Perhaps we can add an option to luceneutil to try non-standard quantization options?

Perhaps a Lucene format (experimental!) that takes an arbitrary "compression factor" (in the range (1, 32)) and internally expands the vector + uses an appropriate quantization?

[mccullocht]

This is a very interesting idea, although I agree it may not be something you want to do in Lucene itself. If it was going to happen inside Lucene you would want the query object to participate and perform expansion before search to avoid doing it per segment.

One aspect that might be interesting to play with is the asymmetric quantization of the query and doc vectors that happens in 1 and 2 bit quantization. Would 6-bit quantization perform better as 6144 dims x 1 bit or 3072 dimx x 2 bits than 1536 dims x 4 bit? Storage of the original vector makes this cost prohibitive but there might be ways to work around this.

Re: compression factor. Another way of expression this might be "I would like each vector to be N bytes long", making it purely about storage/memory costs.

[mikemccand]

Very cool @kaivalnp!

luceneutil's knnPerfTest I think writes an ROC graph for the runs? If you attach that we can visualize the well-interpolated results too.

This approach reminds me of the remarkable innovations in DC (electricity) voltage conversion. It used to be horribly inefficient to change one DC voltageto another, or even to convert home AC electricity to DC (wasteful buzzing "wall warts" from a couple decades ago, heh).

But now they are amazingly efficient, so you can get lots of power (watts) from modern USB-C wall charges in more and more compact packages.

At the heart of this conversion is a fast switching circuit that first converts DC to high voltage pulses, at high frequencies (100s of KHz, maybe low MHz), pumps these spiked pulses into another circuit that smooths it back to DC, higher or lower than the original input.

Modern PC motherboards do this really well, all the VRMs ("voltage regulator modules") that set precise, low voltage for CPU and RAM from 12V, 5V and 3.3V rails provided by the power supply. And respond in microseconds if CPUs suddenly turn on and demand spikes by 100s of watts (I have been learning about VRMs lately because of #15662 and all the heat they generate on beast3 -- Claude thought perhaps the VRM cooling fans broke and that was maybe the cause of massive slowdown).

That adorable tiny USB-C power supply is such nice simple API ("charge device over this wire"), yet under that API is such incredible complexity, including "jump up to jump back down" concept you are using here so well for vector search!

Claude Sonnet 4.6 does a good job explaining these innovations: https://claude.ai/share/342e965f-6ba2-4ad4-801b-b16528cef614

[mikemccand]

And +1 to add such a tool in luceneutil -- what a cool way to turn discrete choices into any continuous choice in between.

Is there much noise in the process? If you have enough index + query vectors, and run your conversion N times (each time a different random rotation), and test recall/performance on each, is the variance small? Maybe for exact NN it's very small (and the the approx NN is just the normal HNSW variance..)?

[mikemccand]

@kaivalnp maybe this could also be an option in knnPerfTest.py? You'd ask it for an in-between quantization and it goes off and rotates and caches things on disk (which it already does for exact NN and Lucene indices)...