High-Rate Nested-Lattice Quantized Matrix Multiplication with Small Lookup Tables
Iris Kaplan, Or Ordentlich
TL;DR
This work tackles the IO bottleneck in matrix multiplication for large-model inference by introducing a rate $R$ hierarchical nested-lattice quantization that quantizes each vector in $M$ layers, enabling inner-product decoding via a small LUT of size $2^{2dR/M}$ while preserving distortion close to single-layer nested-lattice codes. The encoder/decoder are constructed through recursive quantization steps $Q^{\circ m}$ with corresponding increments ${g_m}$, forming the constellation ${\mathcal C}_{L,q,M}$ and guaranteeing exact recovery when $Q^{\circ M}(x)=0$, with overload management and dithering to control errors. The paper provides analytic bounds showing the hierarchical quantizer closely matches the performance of Voronoi codes at the same total rate, and empirical results demonstrate negligible distortion loss compared to standard schemes. Overall, the approach enables high-rate quantization with small LUTs, reducing memory bandwidth requirements and enabling faster, cache-friendly inner-product computations for quantized matrix multiplication in DNN/LLM workloads.
Abstract
Recent work have shown that the quantization for matrix multiplication problem can be optimally solved by quantizing each column in each matrix using a nested lattice code, and then multiplying the de-quantized matrices. It was further demonstrated that when product codes of sub-dimension $d$ and rate $R$ are used, the de-quantization and inner product operations can be implemented with querying a lookup table (LUT) of size $2^{2dR}$, but this is only useful when $dR$ is sufficiently small. This in turn limits LUT-based inner product decoding to low-rate quantizers. In this work, we develop a rate $R$ hierarchical nested lattice quantization framework, which quantizes each vector to $M$ layers, and admits LUT-based inner product decoding using an LUT of size $2^{2d\frac{R}{M}}$, allowing for high-rate quantization. We provide analytic bounds on the loss of the developed scheme compared to standard nested lattice quantizers, and also numerically illustrate that this loss is negligible. Thus, our scheme enables to use small LUTs without compromising the overall distortion.