QK-Normed MLA: QK normalization without full key caching

• URL: http://arxiv.org/abs/2606.16310v1
• Date: Mon, 15 Jun 2026 07:16:30 GMT
• Title: QK-Normed MLA: QK normalization without full key caching
• Abstract summary: We show that QK normalization is a practical stabilization option for MLA models without requiring full-key caching.<n>In our 400M runs trained for up to 100B tokens, QK-Normed MLA achieves lower training loss and better downstream accuracy than QK clipping.
• Score: 65.51862319225098
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Query-key (QK) normalization stabilizes attention by controlling the scale of queries and keys before the dot product, but is not immediately compatible with Multi-head Latent Attention (MLA). MLA achieves efficient decoding by caching low-dimensional latent states instead of full keys, whereas post-projection QK RMSNorm appears to require the fully projected key for every cached token. We show this apparent incompatibility is an implementation artifact, not an architectural constraint. RMSNorm decomposes into a static affine weight and a dynamic scalar RMS statistic. The static key-side weight can be absorbed into the MLA query-side projection; the dynamic key statistic reduces to one inverse-RMS scalar per token and KV group. The resulting formulation is exactly equivalent to explicit post-projection QK RMSNorm in exact arithmetic and preserves MLA's latent decode path. In our 400M runs trained for up to 100B tokens, QK-Normed MLA achieves lower training loss and better downstream accuracy than QK clipping, while H800 decode benchmarks show less than 2% latency overhead up to 256k context. These results make QK normalization a practical stabilization option for MLA models without requiring full-key caching.

Related papers

• Beyond KV Reconstruction: Functional Reconstruction for MLA Draft Models in Speculative Decoding [36.446302092919495]
  We formulate MLA draft construction as functional reconstruction rather than cache compression.<n>We evaluate 192 model-converter-backend-method-task configurations spanning four Llama/Qwen draft-target pairs, TransMLA and MHA2MLA, HF and vLLM, and four 200-prompt tasks.
  arXiv  Detail & Related papers  (2026-07-29T10:54:46Z)
• OScaR: The Occam's Razor for Extreme KV Cache Quantization in LLMs and Beyond [50.440302567029654]
  Multi-modal intelligence has made the Key-Value cache a dominant memory bottleneck for efficient deployment.<n>In this work, we revisit the inherent limitations of the per-channel quantization paradigm.<n>We propose OScaR, an accurate and lightweight KV cache compression framework for X-LLMs.
  arXiv  Detail & Related papers  (2026-05-19T10:53:03Z)
• Latent-Condensed Transformer for Efficient Long Context Modeling [60.72493959155964]
  We propose Latent-Condensed Attention (LCA) for large language models.<n>LCA condenses context within MLA's latent space, where representation is disentangled into semantic latent vectors and positional keys.<n>LCA achieves up to 2.5$times$ prefilling speedup and 90% KV cache reduction at 128K context.
  arXiv  Detail & Related papers  (2026-04-14T08:40:31Z)
• Joint Encoding of KV-Cache Blocks for Scalable LLM Serving [3.3230675313521716]
  Existing KV-cache compression methods rely on rigids, disrupt tensor layouts, or require specialized compute.<n>We propose joint encoding of KV-cache blocks, which fuses similar blocks across requests and input chunks into shared representations.<n>This alleviates the KV-cache memory bottleneck, supporting high-concurrency serving without specialized hardware.
  arXiv  Detail & Related papers  (2026-01-06T14:50:58Z)
• Controlling changes to attention logits [21.620742229590885]
  Stability of neural network weights is critical when training transformer models.<n>Applying normalization to queries and keys, known as QK norm', fixes stability issues in practice, but is not always applicable.<n>We show that our cheap intervention allows us to increase the base learning rate of the network, outperform other methods in the MLA setting, and achieve performance competitive with QK norm.
  arXiv  Detail & Related papers  (2025-11-26T13:24:35Z)
• X-EcoMLA: Upcycling Pre-Trained Attention into MLA for Efficient and Extreme KV Compression [30.770661469301544]
  Multi-head latent attention (MLA) is designed to optimize KV cache memory through low-rank key-value joint compression.<n>We show that our proposed method can effectively compress the KV cache while preserving the performance on the benchmarks.
  arXiv  Detail & Related papers  (2025-03-14T06:49:37Z)
• Quantize What Counts: More for Keys, Less for Values [63.51476878610841]
  Large Language Models (LLMs) suffer inference-time memory bottlenecks dominated by the attention Key-Value ( KV) cache.<n>This paper proposes two theorems that anchor mixed-precision KV quantization in the intrinsic geometry of Transformer models.
  arXiv  Detail & Related papers  (2025-02-20T22:24:27Z)
• Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs [92.7279890407059]
  Multi-head Latent Attention (MLA) is an innovative architecture designed to ensure efficient and economical inference.<n>This paper proposes the first data-efficient fine-tuning method for transitioning from Multi-Head Attention to MLA.
  arXiv  Detail & Related papers  (2025-02-20T18:50:42Z)
• TransMLA: Multi-Head Latent Attention Is All You Need [34.38934956358534]
  TransMLA is a framework that seamlessly converts GQA-based models to MLA-based models.<n>By compressing 93% of the KV cache in LLaMA-2-7B, TransMLA achieves a 10.6x inference speedup at an 8K context length.
  arXiv  Detail & Related papers  (2025-02-11T18:20:18Z)
• Multi-matrix Factorization Attention [59.10039136733939]
  We propose Multi-matrix Factorization Attention (MFA) and MFA-Key-Reuse (MFA-KR)<n>MFA enhances model capacity by efficiently scaling up both the number and dimension of attention heads.<n>MFA-KR further reduces memory requirements by repurposing the key cache as value.
  arXiv  Detail & Related papers  (2024-12-26T15:45:45Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.