Reproduce “QK-Normed MLA: QK normalization without full key caching” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Yizhou Han, Yao Zhao, Jun Zhou, Longfei Li, Ruoyu Sun
The paper introduces QK-Normed MLA, an algebraic rewrite that moves fixed RMS-normalization weights into the query projection and caches one inverse-RMS scalar per token and KV group instead of expanded keys. In matched 400M-parameter, 100B-token runs, it stays below QK-Clip in training loss; accuracy averages 46.33 versus 44.75, but results are not frontier-scale.
Key insights
- RMSNorm can be decomposed into static, feature-wise affine weights and a dynamic, scalar Root Mean Square (RMS) statistic, which allows for separate handling.
- The static affine weights of key-side QK RMSNorm can be absorbed into the query projection, shifting a key-dependent operation to the query side during computation.
- Instead of caching full normalized keys, only a small scalar inverse-RMS statistic needs to be stored per token in the cache, preserving MLA's memory efficiency.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.