Role aware quantization boosts memory efficiency of language model attention
RAQ-MLA: Role-Aware Quantization for Multi-Head Latent Attention Caches
Artificial IntelligenceMachine Learning
Summary
Large language models use a technique called Multi-Head Latent Attention (MLA) to focus on different parts of text efficiently. These models store compressed information to reuse during tasks, but saving this information in fewer bits often reduces accuracy. The authors studied how MLA stores these compressed bits and found two different parts have different roles. They designed a new method, RAQ-MLA, which carefully compresses each part based on their role, making the stored data six times smaller without losing much accuracy.
What this means in practice
- •For large language model engineers: Reduce memory use of latent attention caches in transformer models, enabling larger context windows or faster inference at low bit-precision.
- •For edge ai developers: Deploy more efficient quantized transformer models on devices with limited memory, while keeping reasoning and coding task accuracy nearly intact.
Authors
Tianqi Chen, Xunchao Li, Linchuan Xie, Guanghua Yu, Jianchen Zhu, Peisong Wang, Jian Cheng
Abstract
Multi-Head Latent Attention (MLA) has been adopted by recent frontier large language models. Instead of per-head keys and values, MLA caches a compressed KV latent that is up-projected into keys and values at decode time, together with a narrow positional key that carries position information into the attention score. Existing KV-cache quantization methods are designed around the keys and values of standard attention and do not account for what MLA actually caches, which costs accuracy at low bit widths. We analyze the MLA cache both in activation distribution and in how quantization error propagates to the attention output, and show that the two cached tensors play different roles: the latent mainly supplies the values, whereas the positional key dominates the attention score. We therefore propose RAQ-MLA (Role-Aware Quantization for Multi-Head Latent Attention Caches), which fits an invertible transform to each cached tensor against the error that tensor delivers to the attention output: attention-weighted value error for the latent and score perturbation for the positional key. Without training, RAQ-MLA quantizes the full MLA cache to about 2.5 bits per element, a 6.5x increase in theoretical KV-cache capacity, while remaining nearly lossless on reasoning and coding tasks across four MLA models.