Fast structured quantized transforms speed up model memory compression
Fast Quantized Johnson$-$Lindenstrauss Transforms
Machine LearningData Structures and Algorithms
Summary
Large language models use a lot of memory, especially for storing information during their attention process. Existing methods like Quantized Johnson-Lindenstrauss (QJL) reduce this memory by compressing data, but often rely on slow, dense random projections. The authors propose two new fast methods based on structured transforms called randomized Hadamard transforms combined with Gaussian noise, which keep the benefits of compression while running much faster. They also provide mathematical guarantees about the quality and speed of these new methods, suggesting they work well for large-scale model memory compression.
What this means in practice
- •For machine learning engineers: Reduce memory needs and speed up inference in transformer models by compressing key-value caches with fast quantized embeddings.
- •For signal processing teams: Use fast structured quantized transforms to embed high-dimensional signals efficiently while preserving key inner product properties.
Authors
Angelos Korakitis, Vasileios Nakos
Abstract
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(\varepsilon^{-2}\log^2(d/δ))$, while the second achieves $m=O(\varepsilon^{-2}\log(d/δ))$ under a restriction on $m$. Both constructions can be applied in $O(d\log d)$ time.