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.

• The paper introduces a greedy bucket-wise bit allocation method for Product Quantization and Scalar Quantization that assigns more bits to dimensions producing the greatest validation-recall gains while preserving query-time quantization costs.
• The method consistently outperforms uniform allocation at equal memory budgets, delivering up to 8% relative recall improvement for PQ and 18% for SQ, with the largest gains below 1 bit per dimension.
• The results show that learned allocations concentrate bits in leading dimensions, reflecting variance decay in Matryoshka embeddings, while highlighting open challenges in search efficiency, memory alignment, and billion-scale validation.

Motivation and problem statement

Modern retrieval systems store embeddings at billion scale, where even a 3072-dimensional vector implies terabytes of storage, making quantization indispensable. Existing post-training schemes—most prominently Product Quantization (PQ) and Scalar Quantization (SQ)—are largely embedding-agnostic and allocate bits uniformly across dimensions. The paper observes that this uniformity is misaligned with a structural property of many current commercial embedding models: the Matryoshka property (MRL) (Kusupati et al., 2022), whereby leading dimensions carry a disproportionate share of the representational information. The central question is whether a variable bit allocation scheme can outperform uniform allocation at identical compression rates.

Geometric proxy for the MRL property

The authors note that despite MRL's origins in a nested multi-class cross-entropy objective, no strict mathematical or geometric definition of the resulting embeddings exists. Through empirical investigation they identify an exploitable proxy: variance decay across dimension index. Plotting rolling-window (128-dimension) variance per dimension for OpenAI's text-embedding-3-large and Cohere's embed-v4 reveals prominent iso-variance "levels," with early dimensions exhibiting markedly higher variance than trailing ones. This decay explains why truncation preserves quality—concentrated variance in leading dimensions yields an informative coarse representation without trailing coordinates. Importantly, the allocation algorithm proposed does not use this property explicitly; it emerges from data-driven optimization, providing indirect evidence that variance structure drives the gains.

Greedy bit allocation framework

The framework partitions the $D$-dimensional embedding into $K$ contiguous buckets $B_1,\dots,B_K$ and searches for a per-bucket byte budget vector $\vec{b}$. Starting from a uniform initialization, each iteration evaluates the marginal recall improvement of adding a fixed byte increment $\delta$ to each bucket in turn, measured as 100-recall@100 on a validation set using exact L2 distances over dequantized vectors; the best bucket permanently receives the increment. Runtime scales linearly with $K$ times the number of increments.

Two instantiations are provided:

• Greedy PQ: a bucket receiving $b_k$ bytes is split into exactly $b_k$ contiguous subvectors with 256-center $k$-means codebooks (one byte per centroid index), with remainder dimensions front-loaded into earlier subvectors. Codebooks are cached per bucket-budget pair to keep construction tractable.
• Greedy SQ: allowed widths are $\{0,2,4,8\}$ bits (1-bit is explicitly avoided because it splits the distribution at its mode and distorts unimodal coordinate distributions; 0 bits prunes). Within a bucket, a baseline width $K$0 is applied uniformly and residual budget upgrades $K$1 dimensions to $K$2 via stride $K$3, avoiding metadata overhead while handling sub-1-bit regimes where not every dimension can receive even 2 bits.

The adaptivity is entirely offline at index construction; query-time cost matches static uniform-width quantizers.

Experimental setup

Experiments use DiskANN's quantization utilities (no graph index; brute-force L2 search isolates pure quantization error) on six BEIR datasets—MS Marco, DBpedia-Entity, Quora, FiQA, SciDocs, SciFact—with large corpora truncated to 500k vectors. Embeddings come from text-embedding-3-large (3072d) and embed-v4 (1536d). Budgets sweep from ~0.17 to 1 bit per dimension (bpd) in 1/12 bpd steps; embeddings are split into 8 equal buckets (384 or 192 dimensions each), with greedy sweeps of 40 iterations reaching the 1 bpd target.

Results

Across all dataset–model–quantizer combinations, variable allocation consistently beats uniform allocation at identical budgets, though gains shrink at higher bpd where baseline recall is already high. Headline figures: up to 8% relative recall improvement for PQ and up to 18% for SQ at matched compression rates. Gains are largest in the sub-1-bit regime, where uniform allocation is demonstrably inefficient for MRL embeddings—the setting most relevant to extreme-compression deployments.

The discovered allocations concentrate bits on leading buckets, mirroring the MRL structure, with asymmetry more pronounced for SQ than PQ (which shows smoother decay); the most asymmetric allocations coincide with the largest relative gains. Gains tend to be larger for OpenAI embeddings, consistent with their more distinct iso-variance levels. A comparison against simple leading-dimension truncation at equivalent byte budgets contextualizes absolute recalls—for example, Quora with OpenAI embeddings reaches roughly 51 recall@100 at 1 bpd under truncation—but the paper does not report head-to-head numbers between variable quantization and truncation in the same format, leaving that comparison implicit.

An implication of these results is that structure-aware offline bit allocation is essentially free accuracy relative to uniform allocation, since query-time behavior is unchanged; only index-construction cost increases.

Limitations and open questions

The authors frame the work as preliminary and exploratory, conceding several limitations directly:

• No formal characterization of MRL: whether a mathematical/geometric description derivable from the nested loss exists remains open; such a theory could yield an allocation scheme that exploits MRL by design rather than discovery.
• Inefficient allocation search: the greedy procedure is extremely expensive and impractical at scale; finer-grained buckets might improve quality but would further inflate search complexity. An efficient deterministic heuristic or lightweight randomized alternative is needed.
• Systems challenges: variable-width codes forfeit the memory-alignment and SIMD friendliness of uniform types, posing latency-sensitive engineering problems the paper does not solve.
• Incomplete ablations: sensitivity to bucket count, alternative recall thresholds, hyperparameter tuning, and additional models remains unexplored.

Additionally, evaluation is limited to two proprietary embedding models and corpora of at most 500k vectors, so behavior at billion scale and on non-MRL embeddings is not established.

Conclusion

The paper demonstrates that replacing uniform bit allocation with a greedy, bucket-wise variable scheme yields consistent recall improvements for both PQ and SQ on Matryoshka embeddings at fixed memory budgets, with the largest gains (up to 8% and 18%, respectively) concentrated in the low-bit regime where uniform allocation is least efficient. The recovered allocations independently rediscover the leading-dimension bias implied by MRL, suggesting that variance decay is a usable signal for structure-aware compression. The principal open problems are efficient allocation algorithms grounded in a formal account of the MRL property, and systems-level support for non-uniform code layouts.