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

Token decoder

Open inside →×

learned

3 transformer blocks and a linear readout that turn the reconstructed features into scores over the 50,259 tokens.

L=y3WoutL=y_{3}W_{\mathrm{out}}
All 17 steps
Token-decoder block:\text{Token-decoder block:}
yattn=γ1⊙yn−μ(yn)σ2(yn)+εy_{\mathrm{attn}}=\gamma_{1}\odot\frac{y_n-\mu(y_n)}{\sqrt{\sigma^2(y_n)+\varepsilon}}
ε=10−5\varepsilon=10^{-5}
[ q  k  v ]=yattn WQKV[\,{\htmlClass{vl-stream-q}{q}}\;{\htmlClass{vl-stream-k}{k}}\;{\htmlClass{vl-stream-v}{v}}\,]=y_{\mathrm{attn}}\,W_{QKV}
Rt,j=(cos⁡tθj,  sin⁡tθj){\htmlClass{vl-stream-context}{R_{t,j}}}=\big(\cos t\theta_j,\;\sin t\theta_j\big)
θj=10,000−2j/64\theta_j=10,000^{-2j/64}
Qt=rot⁡ ⁣(Qth; t θ){\htmlClass{vl-stream-q}{Q_{t}}}=\operatorname{rot}\!\big({\htmlClass{vl-stream-q}{Q^{h}_{t}}};\,t\,\theta\big)
θj=10,000−2j/64\theta_j=10,000^{-2j/64}
Kt=rot⁡ ⁣(Kth; t θ){\htmlClass{vl-stream-k}{K_{t}}}=\operatorname{rot}\!\big({\htmlClass{vl-stream-k}{K^{h}_{t}}};\,t\,\theta\big)
θj=10,000−2j/64\theta_j=10,000^{-2j/64}
o=reshape⁡(O)o=\operatorname{reshape}(O)
Δattn=oWO\Delta_{\mathrm{attn}}=oW_{O}
yn′=yn+Δattny'_n=y_n+\Delta_{\mathrm{attn}}
ymlp=γ2⊙yn′−μ(yn′)σ2(yn′)+εy_{\mathrm{mlp}}=\gamma_{2}\odot\frac{y'_n-\mu(y'_n)}{\sqrt{\sigma^2(y'_n)+\varepsilon}}
ε=10−5\varepsilon=10^{-5}
g=ymlpWgateg=y_{\mathrm{mlp}}W_{\mathrm{gate}}
u=ymlpWupu=y_{\mathrm{mlp}}W_{\mathrm{up}}
a=g σ(g)a=g\,\sigma(g)
h=a⊙uh=a\odot u
Δmlp=hWdown\Delta_{\mathrm{mlp}}=hW_{\mathrm{down}}
yn+1=yn′+Δmlpy_{n+1}=y'_n+\Delta_{\mathrm{mlp}}
L=y3WoutL=y_{3}W_{\mathrm{out}}
in
h^\hat h [B batch, T positions, 768 d]
out
LL [B batch, T positions, 50259 vocabulary]
parameters
66,915,072
Parameters

66,915,072 parameters per instance, across its learned components

Exact per-instance accounting
ComponentCalculationParameters
Token-decoder block9,438,720×39{,}438{,}720 \times 328,316,160
Vocabulary readout768×50,259768 \times 50{,}25938,598,912
Per-instance total66,915,072

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.

Explanation

There is no normalization between the last block and the readout.

The readout is a matrix of its own: it is tied neither to GPT-2’s token table nor to anything else.

Sources4

Direct evidence

The token decoder: a rotary table, the configured number of blocks, and a linear readout to the vocabulary without bias.

The implementation, pinned at revision e18c42acdfb7, not yet public

Direct evidence

The token decoder’s blocks: no softcap, no time conditioning, the gated MLP with the configured multiplier.

The implementation, pinned at revision e18c42acdfb7, not yet public

Direct evidence

The latent decoder: 12 heads of 64 features, query and key normalization, softcap 50, 6 layers, MLP multiplier 4; the token decoder has 3 blocks.

The implementation, pinned at revision e18c42acdfb7, not yet public

Direct evidence

The feature model is GPT-2, frozen; the hidden state at index −3 is extracted; the encoder and the decoder are of the Perceiver kind; the file’s vocabulary size and mask token ID are defaults that the trainer overwrites with the tokenizer’s; there are as many latents as positions, of 768 features.

The implementation, pinned at revision e18c42acdfb7, not yet public