Query heads

Pinned implementation and paper LDLM @ e18c42a, not yet public · arXiv:2605.07933v1 · OpenWebText · 1024 · checkpoint, tokenizer and saved statistics not verified · How to read

Latent decoder › 6 decoder blocks › Perceiver attention › Query heads

Query RMSNorm

×

learned

Rescale each head’s 64 features by their root mean square and apply a learned gain.

Q=γQ⊙Qhmax⁡ ⁣(∥Qh∥2/64,  ε){\htmlData{light=queries q-norm}{{\htmlClass{vl-stream-q}{Q}}}}=\gamma_{{\htmlClass{vl-stream-q}{Q}}}\odot\frac{{\htmlData{light=q-heads}{{\htmlClass{vl-stream-q}{Q^{h}}}}}}{\max\!\big(\lVert {\htmlData{light=q-heads}{{\htmlClass{vl-stream-q}{Q^{h}}}}}\rVert_2/\sqrt{64},\;\varepsilon\big)}
ε=10−5\varepsilon=10^{-5}
in
QhQ^{h} [B batch, 12 heads, T positions, 64 head features]
out
QQ [B batch, 12 heads, T positions, 64 head features]
parameters
64 · × 6 blocks = 384

ε bounds the divisor from below; it is not added under the root.

Parameters

64 parameters owned by this component

Exact per-instance accounting
ComponentCalculationParameters
RMSNorm gain646464
Per-instance total64

Across 6 independent blocks: 64×6=38464 \times 6 = 384 parameters.

Derived from the pinned implementation’s constructors and configuration. The checkpoint’s tensors have not been counted. 306,011,162 of these parameters are trained; frozen GPT-2 and the decoder’s mask vector, 124,442,112, are not.

One gain for each of the 64 features.

Sources3

Direct evidence

RMSNorm divides by the vector’s norm times width^(−1/2), clamped from below at ε, and multiplies by a learned gain.

The implementation, pinned at revision e18c42acdfb7, not yet public

Direct evidence

Perceiver attention owns an RMSNorm for each of its two streams, a fused query–key–value projection of the query stream, a fused key–value projection of the other stream, RMSNorms for each head’s queries and keys, and an output projection; no bias.

The implementation, pinned at revision e18c42acdfb7, not yet public

Direct evidence

Both streams are normalized; queries come from one stream, keys and values from both, concatenated along positions; queries and keys are normalized per head; attention runs over all positions with softcap; the heads are merged and projected.

The implementation, pinned at revision e18c42acdfb7, not yet public