Key and value 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 › Key and value heads

Join keys along positions

×

no parameters

Set the two streams’ keys one after the other: 2T key positions.

k=[ kb;  ka ]{\htmlData{light=key-concat}{{\htmlClass{vl-stream-k}{k}}}}=[\,{\htmlData{light=take-k-context}{{\htmlClass{vl-stream-k}{k^{b}}}}};\;{\htmlData{light=take-k-self}{{\htmlClass{vl-stream-k}{k^{a}}}}}\,]
in
kbk^{b} [B batch, T positions, 768 d]kak^{a} [B batch, T positions, 768 d]
out
kk [B batch, 2T positions of both streams, 768 features]
Explanation

A query therefore attends to the latents being decoded and to the decoder’s own state in one softmax.

Sources3

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

Direct evidence

An encoder block adds the Perceiver attention’s output to the latents, then adds the MLP’s output.

The implementation, pinned at revision e18c42acdfb7, not yet public

Direct evidence

A decoder block: Perceiver attention with the queries’ stream first read from the decoder’s state and keys and values from the latents and the state, a scaled residual, an MLP and a second scaled residual.

The implementation, pinned at revision e18c42acdfb7, not yet public