Quantization Beyond Uniform Bit Allocation

Abstract: Quantization is a fundamental technique to handle the growing sizes of embeddings generated by modern models. Existing quantization schemes are largely embedding agnostic and allocate bits uniformly across dimensions. However, recent models produce embeddings with significant geometric structure. In this work, we investigate whether a variable bit allocation scheme can improve quantization quality under a fixed memory budget. We propose a simple variable bit allocation framework that partitions an embedding into contiguous buckets and allocates storage non-uniformly across them. Using a greedy allocation strategy, we instantiate this framework for both Product Quantization (PQ) and Scalar Quantization (SQ). We perform a series of experiments on embeddings known to have the Matryoshka property (MRL), and consistently observe that non-uniform allocations outperform uniform baselines at identical storage budgets. The largest improvements occur in the low-bit regime, where uniform allocation is particularly inefficient for MRL embeddings. At the same compression rates, variable allocation improves recall by up to 8% for PQ and up to 18% for SQ. Our results suggest a new direction for structure-aware compression and indexing techniques for large-scale retrieval systems.

1. Introduction

Nearest neighbor search is a crucial component of nearly all forms of information retrieval in large ML systems that use embeddings—from search to recommendations. However, due to the curse of dimensionality, nearest neighbor search becomes challenging at scale. Modern embeddings routinely have thousands of dimensions; a single 3000-dimensional vector has a disk footprint of 12 kilobytes, implying a total footprint of several terabytes for billion-scale datasets.

Quantization techniques aim to reduce the storage and memory footprints by replacing each full vector with a quantized code that approximates the geometric structure of the dataset. Of these, Product Quantization (PQ) and Scalar Quantization (SQ) are among the most well-known and widely used in large-scale industrial systems.

Furthermore, many modern embeddings (e.g., OpenAI's text-embedding-3-large, Cohere's embed-v4, Gemini Embedding 2, Voyage, Qwen, Jina) exhibit the Matryoshka property (MRL), where the leading dimensions capture a disproportionate fraction of the structure of the embeddings. Prior work observes that across many real-world datasets, retaining just the top 15% of dimensions achieves over 70% recall.

So far, post-training quantization schemes have been largely embedding agnostic. We ask the following question:

"Can a variable bit allocation scheme allow for better quantization of vectors at the same compression rate?"

Specifically, in the setting of MRL embeddings, given the outsized importance of the initial dimensions, a strategy that assigns more bits to the initial dimensions may outperform existing techniques that provide uniform bit budgets to all dimensions.

Through empirical investigation, we identify an exploitable geometric proxy: the decay in variance across dimensions as the dimension index increases. Concentrating variance in the initial dimensions mimics the leading digits of a numerical value, creating a standalone, coarser representation without the trailing dimensions.

Observed variance differences by dimension in OpenAI and Cohere embeddings — **Figure 1:** Observed variance differences by dimension in embeddings produced by OpenAI's `text-embedding-3-large` and Cohere's `embed-v4`. Rolling window average variance per dimension across different datasets (rolling window of 128 dimensions in length) exhibits prominent iso-variance "levels".

2. A Greedy Bit Allocation Scheme

MRL optimizes a multi-class softmax cross-entropy loss $\mathcal{L}: \mathbb{R}^L \times [L] \to \mathbb{R}_+$ across a set of nested dimensions $m \in \mathcal{M}$:

\min_{\{\mathbf{W}^{(m)}\}_{m \in \mathcal{M}}, \, \theta_F} \frac{1}{N} \sum_{i \in [N]} \sum_{m \in \mathcal{M}} c_m \cdot \mathcal{L}\left(\mathbf{W}^{(m)} \cdot F(x_i; \theta_F)_{1:m} \;;\; y_i\right)

Here $F(x_i; \theta_F)_{1:m}$ denotes the truncation of the extracted feature to its first $m$ coordinates, and $c_m \ge 0$ denotes the relative importance scale of dimension $m$ (set to $c_m = 1, \forall m \in \mathcal{M}$ by default).

To evaluate the hypothesis, we propose a variable-bit allocation scheme that partitions the total embedding dimensionality $D$ into $K$ contiguous buckets $B_1, B_2, \dots, B_K$, where bucket $B_k$ has size $d_k = |B_k|$ and $\sum_{k=1}^K d_k = D$.

The greedy search initializes with a uniform byte allocation vector $\vec{b}^{(0)} = (b_{\text{init}}, \dots, b_{\text{init}})$. At each iteration, the algorithm considers candidate configurations by temporarily increasing the bit budget of each bucket in turn by a fixed step-size increment $\delta$ (in bytes), evaluating the search recall $R(\vec{b}, \mathcal{V})$ on a validation query set $\mathcal{V}$ using exact distance computation over the quantized dataset. The bucket yielding the maximum empirical recall receives the $\delta$ bits, repeating until the target budget is achieved.

Algorithm 1: Greedy Bit Allocation

Input: Initial uniform allocation $\vec{b}^{(0)} = (b_{\text{init}}, \dots, b_{\text{init}})$, step-size increment $\delta$, validation set $\mathcal{V}$
Output: Optimized allocation vector $\vec{b}$

1: $\vec{b} \gets \vec{b}^{(0)}$
2: repeat
3: $k^* \gets \arg\max_{k \in \{1, \dots, K\}} R(\vec{b} + \delta \vec{e}_k, \mathcal{V})$
4: $\vec{b} \gets \vec{b} + \delta \vec{e}_{k^*}$
5: until target budget is achieved
6: return $\vec{b}$

Note: The algorithm's runtime is linearly proportional to $K$ times the number of increments needed to reach the target budget. While our algorithm does not explicitly make use of the Matryoshka property in its objective, the resulting bit allocation scheme strongly reflects it, concentrating bits in the leading dimensions.

Greedy Product Quantization

In standard PQ, high-dimensional vector space is decomposed into lower-dimensional, mutually orthogonal subspaces quantized independently using separate $k$-means codebooks. In our variable-bit paradigm, when a bucket $B_k$ of size $d_k$ receives a budget of $b_k$ bytes, it is partitioned into exactly $b_k$ contiguous subvectors. With $C = 256$ centroids, the nearest centroid index fits into $\log_2(256) = 8$ bits ($1$ byte).

Remainder dimensions $r = d_k \bmod b_k$ are front-loaded: the first $r$ subvectors have dimension $s+1$ (where $s = \lfloor d_k / b_k \rfloor$) and the remaining subvectors have dimension $s$. Codebooks $\mathcal{C}_{k}(b)$ are trained once and cached.

Algorithm 2: PQ Dimension Partitioning

Input: Sub-vector $x^{(k)} \in \mathbb{R}^{d_k}$ of bucket $B_k$, budget $b_k$ (in bytes), bucket training data $\mathcal{X}_{\text{train}}^{(k)}$
Output: Quantized code vector $\vec{q}^{(k)} \in \{0, \dots, 255\}^{b_k}$

1: $s \gets \lfloor d_k / b_k \rfloor$, $r \gets d_k \bmod b_k$
2: Partition $x^{(k)}$ and $\mathcal{X}_{\text{train}}^{(k)}$ into $b_k$ contiguous subvectors, where the first $r$ subvectors have dimension $s+1$ and the remaining have dimension $s$.
3: for each subvector index $j = 1$ to $b_k$ do
4: Train codebook $\mathcal{C}_{k, j}$ with 256 centers using $k$-means on the $j$-th partition of $\mathcal{X}_{\text{train}}^{(k)}$
5: Quantize $x^{(k)}_j$ to closest centroid: $\vec{q}^{(k)}[j] \gets \arg\min_{c \in \mathcal{C}_{k, j}} \| x^{(k)}_j - c \|_2$
6: end for
7: return $\vec{q}^{(k)}$

Greedy Scalar Quantization

Scalar Quantization discretizes each dimension independently. Coordinate bit-widths $w_i$ are selected from $\{0, 2, 4, 8\}$. ($1$-bit quantization splits the distribution directly at its mode, artificially pushing high-density central values to the extremes; a bit-width of $0$ corresponds to discarding or pruning the coordinate).

Given a budget of $N_{\text{bits}} = 8 b_k$ for bucket of size $d_k$, we establish a baseline bit-width $W_{\text{base}} \in \{0, 2, 4\}$ and upgrade $u$ dimensions to $W_{\text{next}} \in \{2, 4, 8\}$ using a uniform stride $\Delta = d_k / u$.

Algorithm 3: SQ Bit-Width Allocation

Input: Sub-vector $x^{(k)} \in \mathbb{R}^{d_k}$ of bucket $B_k$, budget $b_k$ (in bytes)
Output: Quantized code vector $\vec{q}^{(k)}$

1: Define allowed bit-widths: $S \gets \{0, 2, 4, 8\}$
2: $N_{\text{bits}} \gets 8 \cdot b_k$
3: $W_{\text{base}} \gets \max \{ w \in S \mid w \cdot d_k \le N_{\text{bits}} \}$
4: if $W_{\text{base}} = 8$ then $W_{\text{next}} \gets 8$, $u \gets 0$
5: else $W_{\text{next}} \gets \min \{ w \in S \mid w > W_{\text{base}} \}$, $u \gets (N_{\text{bits}} - d_k \cdot W_{\text{base}}) / (W_{\text{next}} - W_{\text{base}})$
6: end if
7: $w_i \gets W_{\text{base}}$ \quad for $i = 1, \dots, d_k$
8: if $u > 0$ then
9: $\Delta \gets d_k / u$
10: Upgrade dimensions along the stride: $w_{\text{round}(j \cdot \Delta) + 1} \gets W_{\text{next}}$ \quad for $j = 0, \dots, u-1$
11: end if
12: for each dimension $i = 1$ to $d_k$ do $\vec{q}^{(k)}[i] \gets \text{ScalarQuantize}(x^{(k)}_i, w_i)$ end for
13: return $\vec{q}^{(k)}$

**Figure 2:** Optimized bit allocation across buckets at different bit budgets, noted by bits per dimension (bpd). Observe that optimized allocation concentrates bits towards the initial buckets, in line with what we expect due to the MRL property.

3. Experiments

System: All experiments were run on a Linux x86_64 machine with 16 hardware threads on an Intel Xeon Platinum 8272CL CPU at 2.60 GHz and 32 GiB of main memory, with no GPU. DiskANN was built in C++17 without graph indexing; brute-force L2 distance computation was used to measure exact Nearest Neighbors at all instances.

Datasets: Evaluations were conducted on BEIR benchmark datasets. For MS MARCO, DBpedia-Entity, and Quora, the base corpus is truncated to the first 500,000 entries; smaller datasets (FiQA, SciDocs, SciFact) are used in their entirety.

Dataset	Base Vectors	Queries	Split
MS Marco	500,000	6,980	Dev
DBpedia-Entity	500,000	400	Test
Quora	500,000	10,000	Test
FiQA	57,638	648	Test
SciDocs	25,657	1,000	Test
SciFact	5,183	300	Test

Embedding Models: OpenAI's text-embedding-3-large ($D = 3072$) and Cohere's embed-v4 ($D = 1536$).

Metric: Average $100$-recall@$100$ measured via brute-force exact L2 distance search over dequantized (inflated) database vectors to isolate the quantization error from approximate graph indexing.

Hyperparameters: Dimensionality is partitioned into 8 equal buckets (384 dims/bucket for OpenAI, 192 dims/bucket for Cohere). Search starts at 64 bytes (8 B/bucket) for OpenAI and 32 bytes (4 B/bucket) for Cohere, incrementing by $\delta = 8$ bytes and $\delta = 4$ bytes across 40 iterations up to 1 bit per dimension (bpd).

Comprehensive gains in recall across different datasets — **Figure 3:** A comprehensive account of the gains in recall across different datasets. We compare gains on embeddings generated using OpenAI's `text-embedding-3-large` and Cohere's `embed-v4` models, using both Product Quantization (PQ) and Scalar Quantization (SQ). In addition to the relative recall improvement, we also present the absolute 100-recall@100 achieved by the uniform and variable allocation schemes.

4. Limitations and Future Work

We present this work as a preliminary, exploratory exposition on variable allocation schemes for Matryoshka embeddings. Our results show that it is possible to improve over the uniform bit allocation process using a variable scheme. Important open directions include:

A concrete mathematical description of MRL embeddings: The variance observations in Figure 1 show that variance is a good proxy for the MRL property. However, it remains to show whether a concrete mathematical or geometric description of the embeddings can be derived from the loss function, which would help formalize a quantization scheme that explicitly exploits the MRL property.
An efficient allocation algorithm: The current greedy approach is inefficient and especially slow on large datasets as it iterates over candidate allocations. A practical quantizer requires a more efficient implementation, potentially leveraging a deterministic heuristic, a lightweight randomized/iterative process, or exploring a closed-form expression for bit allocation under the Matryoshka property.
Systems challenges: Dynamic bit allocation introduces systems challenges. Storing large volumes of data using custom-sized data types requires adapting existing architectures, sacrificing the memory access efficiency provided by uniform data types—a particular challenge for SIMD-friendly architectures under latency constraints.
Ablation studies: Thorough ablation studies with different numbers of buckets, further metrics (such as recall at different thresholds), and further hyperparameter tuning across additional models.

5. Acknowledgements

K. S. Sreeramji was supported by the Walmart Centre for Tech Excellence at IISc.