QK-Norm as a Load-Bearing Mechanism in OMN

OMN applies learned RMS normalization independently to query and key representations before rotary position encoding: Q projection → Q RMSNorm → RoPE, and K projection → K RMSNorm → RoPE, while V passes through no such path. Each normalization operates across the 128-dimensional head representation, and the learned gain vectors are shared across the corresponding heads of each layer.

The audited OMN-1 checkpoint contains 7,168 Q/K gain elements with mean 3.460332, population standard deviation 6.431730, minimum −0.050049, maximum 68.5, and 89.6484% of values above 1.43. But the distribution is strongly depth-dependent: direct inspection gives layer-0 means of Q 25.032 and K 24.737, layer-1 means of Q 26.249 and K 26.087, against layer-27 means of Q 1.933 and K 1.910. This directly contradicts older claims that Q/K gains increased with depth.

The historical OMN intervention record reports a strong loss of predictive quality when every Q/K normalization gain was replaced with uniform constants — perplexity moving from 12.97 toward 24.08, 70.63, and 692.08 at constants 2.0, 1.5, and 1.0. But no raw result artifact for those runs is preserved locally, and the evaluation panel behind them — up to fifty 2,048-token sequences from early sorted shards, EOS stripped, no packing metadata — was an internal signal, not a held-out benchmark. The learned gain structure is therefore directly verified, the intervention values remain historically reported, and the causal interpretation remains open. The evidence suggests QK-Norm became deeply integrated into OMN-1's trained attention dynamics — without yet establishing whether that dependence reflects stabilization, compensation, optimization coupling, or another learned mechanism.

Attention logits depend on the interaction between query and key vectors, so changes in their norm can alter logit scale even when relative directional structure stays similar. In OMN's scaled dot-product attention, each logit is a 128-dimensional inner product under the default 1/√d scaling — and with four query heads sharing every KV head through explicit replication, whatever scale regime Q and K carry is broadcast across the grouping. OMN introduces explicit learned RMS normalization into Q and K before positional rotation: instead of leaving query/key magnitude entirely to upstream representation dynamics, the attention pathway gets its own learned scale-regulating mechanism.

No uniqueness is claimed — QK-Norm was not invented at AETHR, and the report does not suggest otherwise. The research contribution is narrower and, arguably, more interesting: how the mechanism behaved inside OMN, how strongly its learned state evolved, and what the intervention appears to reveal.

From the normalized block input u_l: Q, K, and V projections; Q RMSNorm and K RMSNorm; RoPE on Q and K; grouped-query attention across 16 query heads and 4 KV heads; output projection. QK-Norm occurs before RoPE — any diagram showing RoPE feeding normalization is incorrect — and V does not pass through the QK-Norm path at all, consistent with the design reading that values carry content while queries and keys carry alignment geometry. Downstream, four query heads share every KV head through explicit fourfold replication, so the normalized Q/K scale regime broadcasts across the grouping rather than applying per-pair. The implementation applies the same ordering on every forward pass regardless of execution branch: the packed flash-attention path, the per-document fallback, and the masked fallback all consume Q and K after normalization and rotation, so the placement claim holds across all three attention kernels.

For a head representation x ∈ R¹²⁸, the operation measures RMS magnitude with the implementation epsilon of 1e-6, divides it out, and multiplies coordinate-wise by a learned gain vector g:

The nuance that makes the gains worth inspecting: RMS normalization removes absolute scale, but the learned gain then reintroduces a trainable coordinate-wise scale. Every gain vector starts life as 128 ones — the implementation initializes all normalization weights to unity — and normalizes in float32 precision before casting back to the working dtype. Against that starting point, the endpoint is stark: a population mean 3.46× the initialization, five coordinates driven below zero, and a maximum 68.5× above it. QK-Norm is therefore not “forcing all queries and keys to unit scale”: it normalizes the incoming vector and lets the network learn a new scale profile, which is exactly what the checkpoint data below records — a departure from 1.0 to a mean of 3.46 with maxima near 68.

Each Q normalization gain vector — 128 values — is shared across all 16 Q heads of its layer; each K gain vector is shared across that layer's 4 K heads — a 4:1 replication the implementation performs explicitly on every forward pass. The arithmetic closes neatly: 28 layers × 2 vectors × 128 coordinates = 7,168 gain elements, the exact population the checkpoint statistics below describe. No attention head learns an independent normalization vector in the audited implementation, and the report describes no such thing. The learned object is layer-specific and coordinate-specific, which is precisely why replacing it with a scalar destroys two kinds of structure at once.

The overall mean alone misleads, because the distribution is strongly nonuniform with large early-layer values — and with details the mean erases entirely. Five gain coordinates sit below zero, meaning the network learned to invert those coordinates' contribution rather than merely scale it. Over five thousand values — 5,233 of 7,168, nearly seventy-three percent — cluster between 1 and 2, close to their unit initialization, while 59 values exceed 32 and the maximum reaches 68.5. The mechanism must not be summarized as “average gain ≈ 3.46” — the depth structure, developed next, is the finding.

Older documentation narrated gains increasing with depth. Direct checkpoint inspection contradicts it: layer 0 shows Q 25.032 and K 24.737, layer 1 shows Q 26.249 and K 26.087, while layer 27 shows Q 1.933 and K 1.910 — and the steepest change in the network sits between layers 1 and 2, where means collapse from ≈26 to ≈2.5 across a single block boundary. The inspected model carries exceptionally large gain magnitudes in its earliest layers and much smaller means by the final layer. No fuller monotonic trend is inferred beyond what the per-layer data supports.

The profile invites hypotheses and the report marks each as one: early-layer attention may require a different scale regime because the first blocks see raw embedding statistics before the residual stream has been shaped; gains may compensate for upstream or downstream transformations, including the Muon-trained projections whose outputs they rescale; magnitude may interact with residual-stream statistics or optimizer dynamics, since the gains themselves train under AdamW while their inputs arrive from Muon matrices; and coordinate structure may matter more than any scalar mean, given five sub-zero coordinates and a tail past 60. What is not stated is any semantic reading — large early-layer gains do not prove the first layers perform more reasoning, and no such meaning is assigned without evidence.

The historical script performs param.fill_(value): every coordinate of every Q and K normalization gain vector is replaced with the same scalar constant, respecting the parameter's dtype so the intervention introduces no precision confound alongside the value change, and a restore path copies saved state back per parameter name — confirming the checkpoint itself as the control condition. This is a constant replacement intervention — not clipping above a threshold, not scaling the learned pattern, not removing QK-Norm, not selecting layers, not retraining, and not a from-scratch comparison. (The historical filename says “clamping”; the operation it implements is replacement, and the report uses the accurate term.)

Status of everything below: historical reported result — not yet independently reproduced. No raw result JSON is preserved locally, so the values are presented exactly as recorded, with ratios computed programmatically (1.86×, 5.45×, 53.36× baseline) and used sparingly. Note the shape of the response: replacing gains averaging ≈3.46 with 2.0 — a modest rescaling — already costs 11 perplexity points, while the step from 1.5 to 1.0 detonates from 70 to 692. Whatever the mechanism, its sensitivity is nonlinear in the constant, which is itself an observation future experiments must explain rather than average away:

The narrow defensible reading: if the historical numbers reproduce, OMN-1's predictions are highly sensitive to replacing the learned Q/K gain pattern with uniform constants — the gain state became integrated into the trained attention system across all 28 layers, not merely in the high-magnitude early blocks. That is the entire basis for calling the mechanism load-bearing, defined here to mean only that changing the learned mechanism after training causes substantial degradation. It does not mean the model could not have trained without QK-Norm; that counterfactual requires a matched from-scratch ablation that has not been run.

A trained network can reorganize around a component, so intervening after convergence can cause severe degradation even where another architecture trained from scratch could compensate — weights downstream of the gains have spent 193B tokens adapting to the exact scale profile the gains provide, and removing it strands that adaptation with no chance to recover. Post-training intervention asks what the trained checkpoint depends on; from-scratch ablation asks whether training without the mechanism would produce a worse model. OMN currently holds stronger evidence for the first question than the second, and the report never trades one answer for the other.

The replacement g → c·1 changes at least two things at once: overall magnitude (the learned ≈3.46 average scale disappears, replaced by constants as low as 1.0) and coordinate structure (128 distinct learned values per vector collapse to one — including five sub-zero coordinates whose sign flips are erased). Layer identity is a third casualty, with values near 25 and near 2 receiving the identical constant, and Q-versus-K asymmetry a fourth, since both roles are overwritten together and never separately. Today's degradation therefore cannot say whether the operative factor is mean gain magnitude, coordinate pattern, layer-specific scale, Q-versus-K asymmetry, or interaction with another learned component. Disentangling that list is the programme's central unresolved question.

The historical script draws up to 50 sequences of 2,048 tokens from early portions of sorted training shards, with EOS removed, token IDs filtered through tokenizer vocabulary-size logic, and no packed-document metadata at evaluation. Two further caveats attach. First, “early portions of sorted shards” is a distribution statement with consequences: whatever the sort order concentrates early is overweighted, and the panel inherits that skew silently. Second, the 2,048-token evaluation window is far shorter than the 32,768-token training rows, so length-dependent effects are invisible by construction. This is not a protected held-out benchmark panel. The perplexity values are therefore useful as an internal intervention signal — never as publication-grade generalization estimates.

The two columns below use fixed meanings. “Supported” means read off the implementation or the checkpoint, or recorded in the intervention procedure — reproducible from the same artifacts. “Not yet established” means requiring an experiment that has not been run, most often the matched from-scratch ablation of §12:

OMN was designed to make architecture itself part of the experimental surface, and QK-Norm is why that posture matters. The mechanism began as a choice — normalize queries and keys before rotation, learn the returning scale — and the trained checkpoint did not leave those gains near a trivial default. It developed a strongly nonuniform structure with unusually large early-layer values, making QK-Norm a learned part of the checkpoint's operating regime rather than a mere implementation detail. The same lens applies prospectively: any future checkpoint can be profiled identically — 7,168 values, per-layer means, tail counts — and compared against this baseline in minutes, which is what turns a one-off observation into a standing diagnostic.

Each Q/K path spans both optimizer regimes: the projection matrix trains under Muon while its normalization gain trains under AdamW, before RoPE and attention — which means the gain learns its values from inputs produced by Muon-updated projections, while the projections' gradients flow back through attention computed under AdamW-learned scales. That closed loop is what makes the interaction potentially load-bearing in both directions, and it is also what makes the open question genuinely open: the report raises — without answering — the open question of how much learned gain behavior reflects this interaction, and marks it as motivation for future experiments rather than a causal claim.

Experiment 9 is what would convert the observation into a deeper causal result; experiments 1–8 are what would make experiments like the historical one interpretable. The matrix below organizes the set:

The historical result lacks raw artifacts, so the reproduction standard is specified explicitly: an immutable held-out evaluation set (never training shards for the primary result), a fixed tokenization revision, the checkpoint hash, the model code fingerprint, an exact intervention manifest, pure cross-entropy scoring, multiple sequence groups with confidence intervals or bootstrap where appropriate, separated Q and K interventions across the full constant sweep rather than a single point, raw JSON results preserved alongside the exact script revision, figures generated from those artifacts by committed code rather than transcribed by hand, and no training-shard evaluation for the primary result. The JSON requirement is not bureaucratic overhead: machine-readable primary data is what lets a later reader recompute every bar from evidence instead of trusting a transcribed table — the exact failure mode that demoted the historical numbers to reported status.

Historical checkpoint SHA-256 fd8c59f3b7452e85d583b4b3c5d49ff04ae196027c9b8944760ad574979cc728, 1,503,200,256 unique parameters in BF16 storage — 283 tensors, with the embedding and head as separate byte-identical entries — averaging metadata declaring steps 734k / 735k / 736k under a uniform average, the declaration itself stored inside the checkpoint envelope alongside source-file records, method tag, and annotation note. The three source checkpoints were not locally available for re-averaging.

Primary files, repo-relative: src/model.py · src/config.py · scripts/qk_norm_clamping_ablation.py · my-model/model.pt. Architecture sources: OMN-Identity-and-Thesis/02-TECHNICAL-ARCHITECTURE.md · OMN-Identity-and-Thesis/07-EVIDENCE-AND-REFERENCES.md. Historical intervention record: OMN-Documentation/Post Pre-Training Analysis/12-PHASE-9-COMPENSATION-CHECK.md, which tabulates the 12.97 → 24.08 / 70.63 / 692.08 progression — and, just as importantly, the independent audit critique that found template-pattern digits elsewhere in the same documentation generation, which is why this report treats the table as reported rather than reproduced and demands raw JSON before any stronger status. Local absolute filesystem paths are not exposed; references are repository-relative.

Each correction below traces to a verified artifact — implementation source, checkpoint statistics, or recorded procedure — and each replaces a specific oversimplification found in project documentation or natural misreading:

A component becomes load-bearing when the system around it has learned to depend on its presence or state — that is the narrow sense in which the term is used here. Dependence accumulated gradually: gains initialized at unity drifted to a mean of 3.46 with maxima near 68, while downstream projections trained against every intermediate value along the way. OMN-1's Q/K normalization gains are not passive constants; they developed a strongly nonuniform structure during training, and the historical record indicates that replacing this learned structure with increasingly smaller uniform constants sharply degraded predictive behavior. If reproduced under a protected protocol, that would strongly evidence the learned QK-Norm state as part of OMN-1's operating regime. The phrase claims nothing about other Transformers — it describes a property of this trained model.

If Q/K normalization consistently develops substantial learned structure in OMN-derived models, it becomes an architectural feature worth preserving, monitoring, and ablating systematically across scales: similar early-layer behavior in Type-S, the same pattern in Type-M, scale-dependent gain regimes, optimizer-routing interactions, domain-specific profiles in specialized models. Concretely, that means tracking per-layer Q/K means across every Darwin training run the way this report tracks them for OMN-1 — early means near 25, late means near 2, tail counts in the dozens — and treating drift in that profile as a first-class training-health signal rather than a curiosity. None of those behaviors are claimed for future Darwin models today — they are the questions this result creates.

QK-Norm began in OMN as an architectural stabilization mechanism. By the end of training, its learned state appears to have become part of the model's internal operating regime — all 7,168 coordinates of it, from the five sub-zero inversions to the maxima near 68, across every one of the 28 layers.