Head-wise RMSNorm in Attention Mechanisms

• Head-wise RMSNorm is a variant of RMSNorm that normalizes per head subvector rather than the full token vector, enabling more granular scaling in transformers.
• It adapts the standard RMSNorm by applying token-wise normalization independently to each query head and, in grouped-query settings, to each KV group.
• This approach improves computational efficiency and stability in multi-head latent attention and quantization-aware approximations, impacting transformer performance.

Head-wise RMSNorm is a form of root mean square normalization in which the normalization axis is a head-local or group-local feature subvector rather than the full hidden width of a token. In attention, the most precise description is token-wise, feature-axis RMS normalization applied independently to each query head and, in grouped-query settings, independently to each KV group. In common implementations, the RMS statistic is per token and per head, while the affine weight is shared across heads and broadcast over the head axis; in Multi-head Latent Attention (MLA), the key-side statistic is per token and per KV group rather than per query head (Han et al., 15 Jun 2026).

1. Axis semantics and formal definition

The defining operation is standard RMS normalization, but with a changed reduction axis. For a vector $z \in \mathbb{R}^d$,

$$\mathrm{RMSNorm}(z) = (z \odot \gamma)\,\alpha(z), \qquad \alpha(z)=\left(\frac{1}{d}\sum_{j=1}^{d} z_j^2+\varepsilon\right)^{-1/2}.$$

Head-wise RMSNorm applies this operation separately to each head subvector instead of once across the full hidden dimension.

If a token representation is reshaped as $x_t \in \mathbb{R}^{H \times d_h}$, a head-wise form uses

$$\rho_{t,h} = \left( \frac{1}{d_h}\sum_{i=1}^{d_h}x_{t,h,i}^2 \right)^{-1/2}, \qquad y_{t,h,i}=\rho_{t,h}\,x_{t,h,i}\,\gamma_{h,i}.$$

More generally, grouped RMSNorm over feature groups $g$ of size $d_g$ uses

$$\rho_{t,g} = \left( \frac{1}{d_g}\sum_{i \in g}x_{t,i}^2 \right)^{-1/2}, \qquad y_{t,i}=\rho_{t,g(i)}x_{t,i}\gamma_i.$$

The only structural change from standard RMSNorm is therefore the reduction axis: instead of reducing over all hidden features per token, the reduction is performed within each head or feature group (McLean et al., 13 Mar 2026).

Setting	RMS statistic	Reduction domain
Full-width RMSNorm	one scalar per token	all hidden features
Head-wise RMSNorm	one scalar per token and head	head width $d_h$
Grouped RMSNorm	one scalar per token and group	group width $d_g$
GQA/MLA key normalization	one scalar per token and KV group	key feature axis of that KV group

This axis distinction matters because “head-wise RMSNorm” is exact only when the normalized object is indexed by heads. In grouped-query attention and MLA, keys are shared across multiple query heads, so the precise key-side description is per-token, per-KV-group normalization rather than per-query-head normalization (Han et al., 15 Jun 2026).

2. Relation to RMSNorm and LayerNorm

RMSNorm was introduced as a simplification of LayerNorm that removes mean-centering and normalizes only by the root mean square of the summed inputs. In the original formulation,

$$\bar a_i = \frac{a_i}{\mathrm{RMS}(\mathbf a)} g_i, \qquad \mathrm{RMS}(\mathbf a)=\sqrt{\frac{1}{n}\sum_{i=1}^n a_i^2}.$$

The paper argues that LayerNorm’s practical benefit comes primarily from re-scaling invariance rather than re-centering invariance, and describes RMSNorm as giving the model re-scaling invariance property and implicit learning rate adaptation ability (Zhang et al., 2019).

This implies a direct conceptual inheritance for head-wise variants. A head-wise RMSNorm keeps the same RMS-based normalization rule and learned gain, but applies them to each head subvector independently. A plausible implication is that the local invariance structure is preserved head by head: each head is normalized with respect to its own scale, while additive shifts are not removed. The 2019 RMSNorm paper is explicit that RMSNorm is not invariant to all re-centering operations, and that when the mean of summed inputs is zero, RMSNorm is exactly equal to LayerNorm (Zhang et al., 2019).

The same paper also introduced partial RMSNorm, or pRMSNorm, in which the RMS is estimated from only the first $$\mathrm{RMSNorm}(z) = (z \odot \gamma)\,\alpha(z), \qquad \alpha(z)=\left(\frac{1}{d}\sum_{j=1}^{d} z_j^2+\varepsilon\right)^{-1/2}.$$0 of the dimensions: $$\mathrm{RMSNorm}(z) = (z \odot \gamma)\,\alpha(z), \qquad \alpha(z)=\left(\frac{1}{d}\sum_{j=1}^{d} z_j^2+\varepsilon\right)^{-1/2}.$$1 It states that this estimator is biased and that gradient instability was observed when the effective sample size was small. For head-wise designs, this caveat is salient because the feature dimension of a single head is already smaller than the full model width (Zhang et al., 2019).

3. Attention-specific formulations

In attention, head-wise RMSNorm most often appears as post-projection query-key normalization. A baseline attention logit can be written as

$$\mathrm{RMSNorm}(z) = (z \odot \gamma)\,\alpha(z), \qquad \alpha(z)=\left(\frac{1}{d}\sum_{j=1}^{d} z_j^2+\varepsilon\right)^{-1/2}.$$2

where $$\mathrm{RMSNorm}(z) = (z \odot \gamma)\,\alpha(z), \qquad \alpha(z)=\left(\frac{1}{d}\sum_{j=1}^{d} z_j^2+\varepsilon\right)^{-1/2}.$$3 is the query head and $$\mathrm{RMSNorm}(z) = (z \odot \gamma)\,\alpha(z), \qquad \alpha(z)=\left(\frac{1}{d}\sum_{j=1}^{d} z_j^2+\varepsilon\right)^{-1/2}.$$4 is the corresponding KV group. With QK normalization, queries and keys are normalized after projection along their per-head feature dimensions: $$\mathrm{RMSNorm}(z) = (z \odot \gamma)\,\alpha(z), \qquad \alpha(z)=\left(\frac{1}{d}\sum_{j=1}^{d} z_j^2+\varepsilon\right)^{-1/2}.$$5 Here the query statistic is per token and per query head, while the key statistic is per token and per KV group. The learned affine weights $$\mathrm{RMSNorm}(z) = (z \odot \gamma)\,\alpha(z), \qquad \alpha(z)=\left(\frac{1}{d}\sum_{j=1}^{d} z_j^2+\varepsilon\right)^{-1/2}.$$6 are shared across heads and broadcast over the head axis (Han et al., 15 Jun 2026).

This point resolves a common terminological ambiguity. In standard multi-head attention, “head-wise QK norm” is a reasonable shorthand because each head carries its own query and key vector. In grouped-query attention and MLA, however, the key side is not indexed by every query head. The more precise statement is that queries are normalized per token and per query head, whereas keys are normalized per token and per KV group, each over the corresponding feature axis (Han et al., 15 Jun 2026).

The same MLA work further states that its practical implementation uses blockwise normalization because each head is split into content and RoPE components,

$$\mathrm{RMSNorm}(z) = (z \odot \gamma)\,\alpha(z), \qquad \alpha(z)=\left(\frac{1}{d}\sum_{j=1}^{d} z_j^2+\varepsilon\right)^{-1/2}.$$7

and the content and RoPE blocks are normalized separately rather than jointly over the concatenated head vector. Appendix B notes that a unified RMS over $$\mathrm{RMSNorm}(z) = (z \odot \gamma)\,\alpha(z), \qquad \alpha(z)=\left(\frac{1}{d}\sum_{j=1}^{d} z_j^2+\varepsilon\right)^{-1/2}.$$8 is still algebraically compatible, but the blockwise scheme is used because joint normalization creates synchronization between the content path and RoPE path, hurting efficient fused implementation (Han et al., 15 Jun 2026).

4. Head-wise QK RMSNorm in Multi-head Latent Attention

The central technical result of “QK-Normed MLA: QK normalization without full key caching” is that post-projection QK RMSNorm is fully compatible with MLA without caching full projected keys. The apparent incompatibility arises because MLA normally caches a low-dimensional latent state $$\mathrm{RMSNorm}(z) = (z \odot \gamma)\,\alpha(z), \qquad \alpha(z)=\left(\frac{1}{d}\sum_{j=1}^{d} z_j^2+\varepsilon\right)^{-1/2}.$$9 rather than explicit keys, whereas post-projection key RMSNorm seems to require the full projected key $x_t \in \mathbb{R}^{H \times d_h}$0 for every cached token (Han et al., 15 Jun 2026).

The paper resolves this by factoring RMSNorm into a static affine component and a dynamic inverse-RMS scalar. For MLA content keys,

$x_t \in \mathbb{R}^{H \times d_h}$1

with $x_t \in \mathbb{R}^{H \times d_h}$2 equal to one inverse-RMS scalar per token and KV group. Because $x_t \in \mathbb{R}^{H \times d_h}$3 is static and independent of token $x_t \in \mathbb{R}^{H \times d_h}$4, it can be folded into the query-side latent projection. Defining

$x_t \in \mathbb{R}^{H \times d_h}$5

the normalized content score becomes

$x_t \in \mathbb{R}^{H \times d_h}$6

In exact arithmetic, this absorbed formulation is identical to explicitly materializing the post-up-projection key and applying QK RMSNorm before the dot product; in floating point, slight differences may occur because the two paths use different multiplication and reduction orders (Han et al., 15 Jun 2026).

The implementation consequence is that MLA keeps its latent decode path. Instead of caching full keys of shape $x_t \in \mathbb{R}^{H \times d_h}$7, it caches the latent states $x_t \in \mathbb{R}^{H \times d_h}$8 and an auxiliary scalar cache $x_t \in \mathbb{R}^{H \times d_h}$9 of shape $$\rho_{t,h} = \left( \frac{1}{d_h}\sum_{i=1}^{d_h}x_{t,h,i}^2 \right)^{-1/2}, \qquad y_{t,h,i}=\rho_{t,h}\,x_{t,h,i}\,\gamma_{h,i}.$$0. The paper explicitly warns against absorbing the key scalar into the cached latent vector because $$\rho_{t,h} = \left( \frac{1}{d_h}\sum_{i=1}^{d_h}x_{t,h,i}^2 \right)^{-1/2}, \qquad y_{t,h,i}=\rho_{t,h}\,x_{t,h,i}\,\gamma_{h,i}.$$1 differs across KV groups and doing so would expand the cache from $$\rho_{t,h} = \left( \frac{1}{d_h}\sum_{i=1}^{d_h}x_{t,h,i}^2 \right)^{-1/2}, \qquad y_{t,h,i}=\rho_{t,h}\,x_{t,h,i}\,\gamma_{h,i}.$$2 to $$\rho_{t,h} = \left( \frac{1}{d_h}\sum_{i=1}^{d_h}x_{t,h,i}^2 \right)^{-1/2}, \qquad y_{t,h,i}=\rho_{t,h}\,x_{t,h,i}\,\gamma_{h,i}.$$3, eliminating much of MLA’s memory advantage (Han et al., 15 Jun 2026).

This result sharpens the meaning of head-wise RMSNorm in latent attention. On the query side, normalization remains per token and per query head. On the key side, the normalization is per token and per KV group, and the normalized key need not be cached explicitly. The method is therefore not an approximation in the architectural sense, but an algebraically exact rewrite of post-projection QK RMSNorm for MLA (Han et al., 15 Jun 2026).

5. Computational variants and quantization-aware approximations

MXNorm is a drop-in replacement for RMSNorm that estimates the RMS using only block scales calculated as part of the MXFP8 cast. Its native setting is standard token-wise RMSNorm over the full hidden dimension $$\rho_{t,h} = \left( \frac{1}{d_h}\sum_{i=1}^{d_h}x_{t,h,i}^2 \right)^{-1/2}, \qquad y_{t,h,i}=\rho_{t,h}\,x_{t,h,i}\,\gamma_{h,i}.$$4, not per-head normalization. For $$\rho_{t,h} = \left( \frac{1}{d_h}\sum_{i=1}^{d_h}x_{t,h,i}^2 \right)^{-1/2}, \qquad y_{t,h,i}=\rho_{t,h}\,x_{t,h,i}\,\gamma_{h,i}.$$5, standard RMSNorm computes one scalar $$\rho_{t,h} = \left( \frac{1}{d_h}\sum_{i=1}^{d_h}x_{t,h,i}^2 \right)^{-1/2}, \qquad y_{t,h,i}=\rho_{t,h}\,x_{t,h,i}\,\gamma_{h,i}.$$6 per token from all hidden features, whereas MXNorm uses per-block absmax summaries and reduces over the resulting block statistics instead of over all $$\rho_{t,h} = \left( \frac{1}{d_h}\sum_{i=1}^{d_h}x_{t,h,i}^2 \right)^{-1/2}, \qquad y_{t,h,i}=\rho_{t,h}\,x_{t,h,i}\,\gamma_{h,i}.$$7 features (McLean et al., 13 Mar 2026).

The core motivation is that in pre-norm transformers with MX quantization, RMSNorm and MX casting both scan activations along the hidden dimension. MXNorm reuses the block absmax values already computed for quantization to approximate the inverse RMS, giving a 32x decrease in the size of the reduction needed for normalization for MXFP8 with block size $$\rho_{t,h} = \left( \frac{1}{d_h}\sum_{i=1}^{d_h}x_{t,h,i}^2 \right)^{-1/2}, \qquad y_{t,h,i}=\rho_{t,h}\,x_{t,h,i}\,\gamma_{h,i}.$$8. The paper reports practical kernel speedups using only torch.compile of up to 2.4x for MXNorm over RMSNorm, corresponding to a 1.3% speedup in Llama 3 8B transformer layers in MXFP8 and a 2.6% speedup in NVFP4 (McLean et al., 13 Mar 2026).

For head-wise RMSNorm, the paper does not provide a direct experiment. It does, however, give a mathematically straightforward grouped analogue: aggregate only over the blocks belonging to one head instead of all blocks in the token. This suggests a head-wise MXNorm could define per-head block maxima $$\rho_{t,h} = \left( \frac{1}{d_h}\sum_{i=1}^{d_h}x_{t,h,i}^2 \right)^{-1/2}, \qquad y_{t,h,i}=\rho_{t,h}\,x_{t,h,i}\,\gamma_{h,i}.$$9 and a per-head inverse-RMS estimate $g$0 using only those block summaries. The paper’s own discussion makes the practical constraint explicit: such reuse is most natural when each head’s feature slice is contiguous and the head size is divisible by the MX block size, so that blocks do not straddle head boundaries (McLean et al., 13 Mar 2026).

The same work also provides a caution about approximation quality. It reports that MXNorm($g$1) matched RMSNorm well at 8B scale, while MXNorm($g$2) underperformed and the “post-round MXNorm” variant, which estimates RMS from rounded E8M0 scales rather than pre-round block absmaxes, was unstable at 8B. A plausible implication for head-wise use is that smaller normalization groups may make approximation fidelity more sensitive to block count and block alignment (McLean et al., 13 Mar 2026).

6. Symmetry, gauge, and coordinate identity

Standard residual-stream RMSNorm with generic per-channel gain induces a native signed-permutation gauge. Theorem 2.1 of “Signed-Permutation Coordinate Transport for RMSNorm Transformers” states that for RMSNorm with generic per-channel gain $g$3, the maximal subgroup of $g$4 that preserves native coordinate identity without absorbing $g$5 into adjacent weights is

$g$6

LayerNorm, by contrast, has only permutation gauge up to one global sign flip. The reason is that LayerNorm’s mean-centering breaks equivariance under independent sign flips, whereas RMSNorm lacks mean-centering and retains them (Sweeney, 30 Jun 2026).

The paper does not explicitly analyze head-wise RMSNorm. A careful extrapolation from its theorem is that a fixed block-wise or head-wise RMSNorm would likely replace the global gauge $g$7 by a product of per-block signed-permutation groups,

$g$8

assuming generic gains within each block and a fixed normalization partition. This suggests that within-head signed permutations would remain native symmetries, while cross-head coordinate mixing would generally cease to be native because it would change which coordinates contribute to which RMS denominator (Sweeney, 30 Jun 2026).

This gauge-theoretic perspective matters for any coordinate-indexed artifact. For standard RMSNorm models, permutation-only alignment is symmetry-incomplete, and the paper’s sign-marginalized Hungarian matching solves assignment on $g$9 rather than raw signed correlation. Its empirical results show why this matters: on RMSNorm models such as Qwen2.5-7B and Llama-3.1-8B, signed-cost matching recovers roughly 50% of coordinates under a sampled $d_g$0 basis change, whereas sign-marginalized matching recovers 100.0%. The same paper reports that TinyLlama SAE reconstruction has NMSE 0.004 under $d_g$1 transport versus 1.08 under $d_g$2, and that Qwen sentiment steering preserves 95.8% of its effect under $d_g$3 versus 17.2% under $d_g$4 (Sweeney, 30 Jun 2026).

For head-wise RMSNorm, the natural implication is local rather than global gauge structure. Coordinate identity would be meaningful relative to an explicit signed gauge within each head or block, and any head-local sparse edit, steering vector, or optimizer state would need to be transported with signs, not just permutations. The paper does not test this configuration directly, but its RMSNorm theory points in that direction (Sweeney, 30 Jun 2026).

7. Empirical behavior, caveats, and scope

The empirical record is strongest for three neighboring objects: RMSNorm itself, post-projection QK RMSNorm in attention, and quantization-aware RMS approximations. The original RMSNorm paper reports that RMSNorm achieves comparable performance against LayerNorm but reduces the running time by 7%~64% on different models. It also reports Transformer results in which the no-normalization baseline failed, LayerNorm achieved Test14 26.6 and Test17 27.7 at 248s/1k steps, and RMSNorm achieved Test14 26.8 and Test17 27.7 at 231s/1k steps (Zhang et al., 2019).

For head-wise QK normalization in latent attention, the most direct evidence comes from QK-Normed MLA. In 400M-parameter MLA models trained for 100B tokens, QK-Normed MLA shows consistently lower training loss than QK-Clip over the entire run, improves 7 of 8 reported accuracy tasks, raises the average score from 44.75 to 46.33, and improves LAMBADA perplexity from 16.28 to 14.18. Mechanism-level diagnostics reported in the same work state that QK-Normed MLA keeps attention logits in a lower range than QK-Clip and shows a smoother gradient norm profile. Under an intentionally high learning rate of $d_g$5, QK-Clip diverges with gradient NaNs at step 884, while QK-Normed MLA remains stable and continues reducing loss. On H800 decode benchmarks, the latency overhead ranges from 0.95% to 1.87% up to 256k context, with average overhead 1.40% and overall overhead under 2% (Han et al., 15 Jun 2026).

For quantization-aware normalization, MXNorm provides evidence that RMS estimation can be made substantially cheaper without large losses in model quality, but not specifically at head granularity. On 8B models trained on 300B tokens, the reported final loss is 2.132 for RMSNorm, 2.175 for MXNorm($d_g$6), and 2.126 for MXNorm($d_g$7); on zero-shot OLMES benchmarks, RMSNorm and MXNorm($d_g$8) each lead on 5 of 10 tasks. The same paper emphasizes two practical caveats that are especially relevant to head-wise variants: approximation quality improves with the number of blocks, and reuse of block statistics is most natural when heads align cleanly with quantization blocks (McLean et al., 13 Mar 2026).

Several misconceptions are clarified by this literature. First, head-wise RMSNorm is not a distinct normalization principle from RMSNorm; it is RMSNorm with a different reduction axis. Second, in grouped-query and MLA settings, “head-wise” is precise for queries but not for keys, whose normalization is per token and per KV group. Third, in MLA, post-projection QK RMSNorm does not require full key caching; the need to cache full keys is an implementation artifact rather than an architectural constraint (Han et al., 15 Jun 2026).

Taken together, these results place head-wise RMSNorm at the intersection of three concerns: normalization geometry, attention parameterization, and systems efficiency. The directly demonstrated cases show that RMS-based scaling can be localized to heads or KV groups in attention, can be algebraically absorbed into latent-attention decode paths without losing equivalence, and can in some settings be approximated from quantization metadata. The broader literature also suggests that once normalization is localized, symmetry and coordinate identity become localized as well, with signed permutations rather than unsigned permutations providing the appropriate native gauge for RMSNorm-based models (Sweeney, 30 Jun 2026).