Architecture Name
MorphNorm with MLA
Parent issue
#1
Motivations
QK-Norm stabilizes Transformer training by controlling query and key scales, but naive post-projection key normalization is incompatible with MLA's absorbed decoding path. QK-Normed MLA resolves this exactly by caching a projected-key inverse-RMS scalar for every token and KV group. This preserves the latent cache, but requires a temporary key up-projection at cache write time, an additional per-group scalar cache, and score-stage scaling. QuacK instead controls logit changes through parameter-dependent learning rates.
MorphNorm explores a different trade-off: approximate the projected-key RMS using the shared latent RMS and a calibrated head-wise gain. The hypothesis is that this statistical approximation retains most of QK-Norm's logit control while avoiding full-key caching, per-token/per-group key-norm caches, and cache-write key materialization. See also this analysis.
Proposed Architecture
Let $c_t=W^{DKV}x_t$ be the shared KV latent and $k^C_{t,h}=W^{UK}_h c_t$ the materialized content key. Define the token-dependent projection gain and its calibrated head-wise estimate as
$$
r_{t,h}=\frac{\mathrm{RMS}(W^{UK}_h c_t)}{\mathrm{RMS}(c_t)}, \qquad s_h=\mathbb{E}_t[r_{t,h}].
$$
Exact projected-key QK-Norm divides by $r_{t,h}\mathrm{RMS}(c_t)$. MorphNorm replaces the dynamic $r_{t,h}$ with $s_h$:
$$
\widehat{k}^{C}_{t,h}=\frac{\gamma\odot W^{UK}_h c_t}{s_h\mathrm{RMS}(c_t)}.
$$
The approximation error relative to exact QK-Norm is the scalar $r_{t,h}/s_h$. MorphNorm is therefore most accurate when the projection gain concentrates across tokens.
The static factor remains absorbable into the query-side projection:
$$
(\widehat{q}^C_h)^\top\widehat{k}^C_{t,h}=\left((\mathrm{diag}(\gamma/s_h)W^{UK}_h)^\top\widehat{q}^C_h\right)^\top\frac{c_t}{\mathrm{RMS}(c_t)}.
$$
Queries use learned RMSNorm, while the small shared RoPE key is normalized independently. The prototype estimates $s_h$ across tokens and data-parallel workers; the inference implementation will use a calibrated and frozen value. The absorbed path keeps the raw latent for the value computation and needs at most one shared inverse-latent-RMS scalar per token for the content-key path, rather than one projected-key scalar per token and KV group.
Preliminary Results (if any)
https://api.wandb.ai/links/antnlp/2irrfeu8
Experiments Plan
All training comparisons will use the Muon optimizer and matched architecture, data, token budget, context length, initialization, learning-rate schedule, and evaluation intervals.
1B validation
Key metrics
Scaling decision
Architecture Name
MorphNorm with MLA
Parent issue
#1
Motivations
QK-Norm stabilizes Transformer training by controlling query and key scales, but naive post-projection key normalization is incompatible with MLA's absorbed decoding path. QK-Normed MLA resolves this exactly by caching a projected-key inverse-RMS scalar for every token and KV group. This preserves the latent cache, but requires a temporary key up-projection at cache write time, an additional per-group scalar cache, and score-stage scaling. QuacK instead controls logit changes through parameter-dependent learning rates.
MorphNorm explores a different trade-off: approximate the projected-key RMS using the shared latent RMS and a calibrated head-wise gain. The hypothesis is that this statistical approximation retains most of QK-Norm's logit control while avoiding full-key caching, per-token/per-group key-norm caches, and cache-write key materialization. See also this analysis.
Proposed Architecture
Let$c_t=W^{DKV}x_t$ be the shared KV latent and $k^C_{t,h}=W^{UK}_h c_t$ the materialized content key. Define the token-dependent projection gain and its calibrated head-wise estimate as
Exact projected-key QK-Norm divides by$r_{t,h}\mathrm{RMS}(c_t)$ . MorphNorm replaces the dynamic $r_{t,h}$ with $s_h$ :
The approximation error relative to exact QK-Norm is the scalar$r_{t,h}/s_h$ . MorphNorm is therefore most accurate when the projection gain concentrates across tokens.
The static factor remains absorbable into the query-side projection:
Queries use learned RMSNorm, while the small shared RoPE key is normalized independently. The prototype estimates$s_h$ across tokens and data-parallel workers; the inference implementation will use a calibrated and frozen value. The absorbed path keeps the raw latent for the value computation and needs at most one shared inverse-latent-RMS scalar per token for the content-key path, rather than one projected-key scalar per token and KV group.
Preliminary Results (if any)
https://api.wandb.ai/links/antnlp/2irrfeu8
Experiments Plan
All training comparisons will use the Muon optimizer and matched architecture, data, token budget, context length, initialization, learning-rate schedule, and evaluation intervals.
1B validation
Key metrics
Scaling decision