[ARCH-PROP] MorphNorm

### 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](https://arxiv.org/abs/2606.16310) 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](https://arxiv.org/abs/2511.21377) 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](https://kexue.fm/archives/11126).

### 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**

- [ ] Compare MLA+MorphNorm against standard MLA, QK-Clip / MuonClip, QuacK for MLA, and exact [QK-Normed MLA](https://arxiv.org/abs/2606.16310) under the same learning-rate sweep and random seeds.

**Key metrics**

- [ ] Record maximum attention logits, loss spikes or divergence, evaluation loss/perplexity, and downstream-task accuracy.
- [ ] Record training throughput, peak memory, decode latency, cache-write cost, and KV-cache size over multiple context lengths.

**Scaling decision**

- [ ] Proceed to 3B and 8B if the 1B results show competitive logit control and model quality while reducing cache or decode cost relative to exact QK-Normed MLA.