Fast Quantized Johnson - Lindenstrauss Transforms

The gist: This paper introduces two constructions, D-FQJL and P-FQJL, for fast asymmetric quantized Johnson–Lindenstrauss transforms that enable O(d log d) time complexity while preserving analogous guarantees to Fast JL in the non-quantized setting.

Introduction and Motivation

Vector quantization (VQ) is a crucial technique for reducing the memory footprint of large language models during inference, particularly by compressing Key-Value (KV) caches. The core idea behind VQ is to keep low-bit representations of high dimensional vectors while retaining quantities of interest like inner products or mean squared errors with minimal error. This is significant because the memory footprint of the KV cache can be a bottleneck in transformer-based architectures, and compressing keys, values, and queries without sacrificing accuracy allows LLMs to process longer context windows. The framework is also applicable to nearest neighbor search in vector databases where searches rely on inner product or cosine similarity measures.

Fast Quantized Johnson–Lindenstrauss Transforms

The paper develops fast structured analogues of the asymmetric quantized Johnson–Lindenstrauss estimator, aiming to replace the dense Gaussian projection with a transform applicable in nearly linear time while retaining the same asymmetric storage format. Both constructions are carefully designed twists of compositions of two Randomized Hadamard Transforms (RHTs) and provide results that are new and incomparable to existing work.

1. D-FQJL: Dithered Fast Quantized JL

The D-FQJL estimator smooths the sign quantizer using non-subtractive Gaussian dithering, where the noise is added to the key sketch before quantization. The estimator is defined as eˆσ(q, k):= p π/2 √1 + σ squared m∥k∥2 ⟨T q, Qσ(k)⟩. Unbiasedness of D-FQJL is proven by showing that for every q ∈ R d, every nonzero k ∈ R d, and every σ > 0, the expectation is Eg[ˆeσ(q, k) I, D±] = ⟨q, k⟩. The distortion theorem requires a sample size m such that m ≳ (1 + σ 2)(1 + 1/σ) squared ε squared log d δ log 1/δ.

2. P-FQJL: Permutation-Based Fast Quantized JL

The P-FQJL estimator retains the original sign quantizer but introduces a uniformly random permutation between the two Hadamard transforms, resulting in TΠ:= SI √d HDg Π HD±. The permutation's role is to obtain concentration, as it randomizes how coordinates are paired with shared Gaussian entries. The covariance regularization by permutation shows that under certain conditions on m, the blocks of the covariance matrix are close to their ideal values, leading to a distortion bound. The distortion theorem requires a sample size m such that m ≥ 1/5ε squared log d δ.

Key-Cache Compression Application

The constructions are applied directly to KV-cache compression, where the key is stored as an m-bit code and one scalar norm, while the query is transformed on demand into a compact representation. The D-FQJL key-cache distortion theorem guarantees that for keys k1,..., kn and a query q such that maxj∥kj∥2 ≤ r and∥q∥2 ≤ r, the error satisfies score [σ,j (n) − scorej (n) ≲ ε scorej (n)]. Similarly, the P-FQJL key-cache distortion theorem guarantees that for keys k1,..., kn and a query q such that maxj∥kj∥2 ≤ r and∥q∥2 ≤ r, the error satisfies score [Π,j (n) − scorej (n) ≲ 20ε scorej (n)].

Connection to TurboQuant

The paper connects these constructions to TurboQuant and earlier DRIVE/EDEN line of work, which use random rotations or RHTs for fast practical implementations. The authors argue that directly replacing the Gaussian projection in the asymmetric sign-linear estimator by HD± does not preserve exact unbiasedness, but their constructions address this gap by retaining fast-transform complexity while recovering an exactly unbiased asymmetric estimator. This provides provably accurate fast replacements for TurboQuant’s dense residual QJL stage.

Space Complexity

The memory cost for both D-FQJL and P-FQJL involves a shared sketching operator that requires O(d+m) storage, as the Hadamard matrices are implicit. Representing a key–query pair requires an m-bit code and one FPN for the key norm, along with an m-dimensional floating-point sketch of the query. Consequently, a cache of n keys requires nm bits and n FPNs in addition to the single shared structured operator.

Conclusion

The work provides running time and embedding dimensions analogous to those in fast Johnson–Lindenstrauss transforms, proving that the power of dimensionality reduction persists under scalar quantization. Both constructions demonstrate that they can be applied directly to KV-cache compression. The paper concludes by providing detailed proofs of the claims stated in the main text.

A Toolkit

The toolkit section defines notation for vectors, norms, inner products, and random variables like Rademacher and standard Gaussian distributions. It also introduces the transforms T:= SI √d HDgHD± and TΠ:= SI √d HDgΠHD±. The probabilistic tools used include Hoeffding’s inequality, Gaussian concentration, and the Matrix Bernstein inequality for random permutations. Lemma 6 provides a Gaussian sign identity that is fundamental to both estimators. Theorem 5 and Theorem 6 provide bounds on independent variables and Lipschitz functions, respectively. The proof of D-FQJL distortion relies on Lemmas 9, 10, and 11 to bound the error terms. The proof of P-FQJL distortion utilizes Lemma 13 to bound the covariance blocks via permutation. Finally, Theorem 8 provides a Gaussian quadratic-form concentration inequality used in the analysis of the unquantized fast JL transform.

Detailed Proofs of D-FQJL and P-FQJL

The detailed proofs for D-FQJL and P-FQJL are provided, including elementary Gaussian identities like Lemma 6. The proof for D-FQJL distortion involves a sequential union bound over events E1, E2, E3, and E4 to show that the error is bounded. The proof for P-FQJL distortion involves a three-stage process: permutation regularization, Gaussian regression separation of correlated components, and Gaussian concentration control. Both proofs rely on the dimension assumption m ≳ (1 + σ 2)(1 + 1/σ) squared ε squared log d δ log 1/δ for D-FQJL and m ≥ 16r 4/5ε 2d log dn δ for P-FQJL. The final results show that the error is bounded by ε√d/4 for D-FQJL and 20ε scorej(n) for P-FQJL.

Conclusion

In summary, the paper provides fast quantized Johnson–Lindenstrauss transforms with corresponding asymmetric estimators, proving that running time and embedding dimensions analogous to those in fast Johnson–Lindenstrauss transforms persist under scalar quantization. The two constructions, D-FQJL and P-FQJL, offer provably accurate fast replacements for TurboQuant’s dense residual QJL stage. Both methods achieve the desired performance guarantees simultaneously for every key and query pair in the KV cache compression setting.

References

Krish Agarwal, Rishi Astra, Adnan Hoque, Mudhakar Srivatsa, Raghu Ganti, Less Wright, and Sijia Chen. Hadacore: Tensor core accelerated hadamard transform kernel. arXiv preprint arXiv:2412.08832, 2024.

Nir Ailon and Bernard Chazelle. Approximate nearest neighbors and the fast johnson–lindenstrauss transform. In Proceedings of the thirty-eighth annual ACM symposium on Theory of computing, pages 557–563, 2006.

Nir Ailon and Bernard Chazelle. The fast johnson–lindenstrauss transform and approximate nearest neighbors. SIAM Journal on computing, 39(1):302–322, 2009.

Alexandr Andoni, Piotr Indyk, Thijs Laarhoven, Ilya Razenshteyn, and Ludwig Schmidt. Practical and optimal lsh for angular distance.

1. Fast KV Cache Compression for LLMs: The D-FQJL construction allows storing each key as an m-bit binary sketch while the query is transformed to a compact representation, enabling KV cache compression that is data-oblivious and requires only O(d+m) space for the sketching operator.

2. Enhanced Attention Mechanism Efficiency: The P-FQJL construction, utilizing a permutation between Hadamard transforms, allows for an embedding in O(d log d) time, providing a way to replace dense projections with structured transforms while retaining the asymmetric structure of QJL.

3. Provably Accurate Residual Modeling: D-FQJL addresses the gap where replacing Gaussian projections with RHTs might bias results by using Gaussian–Hadamard transforms that retain 1-bit quantization while preserving the exact sign–linear expectations required for unbiasedness.

4. Scalable Context Windows: By achieving running time and embedding dimensions analogous to ones in fast Johnson–Lindenstrauss transforms, the system can process much longer context windows and support larger batch sizes without meeting GPU memory barriers.

Quantization has emerged as one of the primary techniques for reducing the memory footprint of large language models during inference. A prominent example is Quantized Johnson - Lindenstrauss (QJL), which compresses KV caches using 1-bit asymmetric quantization while approximately preserving the inner products required by the attention mechanism. Recent theoretical and practical work aims at replacing the dense Gaussian projections used by the aforementioned transforms with more structured embeddings - such as compositions of randomized Hadamard transforms - while retaining the same embedding dimension. In this paper, we introduce two constructions based on a composition of randomized Hadamard transforms, albeit with a Gaussian rather than Rademacher diagonal in the second stage. The first adds Gaussian dithering to smooth the quantizer, while the second inserts a random permutation to regularize the covariance structure of the resulting Gaussian measurements. We prove embedding-dimension and running-time guarantees analogous to those of Fast JL in the non-quantized setting. The first construction achieves m=O(epsilon-2 2(d/δ)), while the second achieves m=O(epsilon-2 (d/δ)) under a restriction on m. Both constructions can be applied in O(d d) time.