LDLM · Token-decoder block

Token decoder

Token-decoder block

learned

One of 3 pre-norm blocks: self-attention with rotary positions and a gated MLP.

y_{\mathrm{mlp}}=\gamma_{2}\odot\frac{y'_n-\mu(y'_n)}{\sqrt{\sigma^2(y'_n)+\varepsilon}}

\varepsilon=10^{-5}

y_{n+1}=y'_n+\Delta_{\mathrm{mlp}}

All 21 steps

y_{\mathrm{attn}}=\gamma_{1}\odot\frac{y_n-\mu(y_n)}{\sqrt{\sigma^2(y_n)+\varepsilon}}

\varepsilon=10^{-5}

\text{Self-attention:}

[\,{\htmlClass{vl-stream-q}{q}}\;{\htmlClass{vl-stream-k}{k}}\;{\htmlClass{vl-stream-v}{v}}\,]=y_{\mathrm{attn}}\,W_{QKV}

{\htmlClass{vl-stream-context}{R_{t,j}}}=\big(\cos t\theta_j,\;\sin t\theta_j\big)

\theta_j=10,000^{-2j/64}

{\htmlClass{vl-stream-q}{Q_{t}}}=\operatorname{rot}\!\big({\htmlClass{vl-stream-q}{Q^{h}_{t}}};\,t\,\theta\big)

\theta_j=10,000^{-2j/64}

{\htmlClass{vl-stream-k}{K_{t}}}=\operatorname{rot}\!\big({\htmlClass{vl-stream-k}{K^{h}_{t}}};\,t\,\theta\big)

\theta_j=10,000^{-2j/64}

S=QK^{\mathsf T}

F=\frac{S}{\sqrt{64}}

A_{t,s}=\frac{\exp F_{t,s}}{\sum_{u=1}^{T}\exp F_{t,u}}

O=AV

o=\operatorname{reshape}(O)

\Delta_{\mathrm{attn}}=oW_{O}

y'_n=y_n+\Delta_{\mathrm{attn}}

y_{\mathrm{mlp}}=\gamma_{2}\odot\frac{y'_n-\mu(y'_n)}{\sqrt{\sigma^2(y'_n)+\varepsilon}}

\varepsilon=10^{-5}

\text{Gated MLP:}

g=y_{\mathrm{mlp}}W_{\mathrm{gate}}

u=y_{\mathrm{mlp}}W_{\mathrm{up}}

a=g\,\sigma(g)

h=a\odot u

\Delta_{\mathrm{mlp}}=hW_{\mathrm{down}}

y_{n+1}=y'_n+\Delta_{\mathrm{mlp}}

in: $\hat h$ [B batch, T positions, 768 d]
out: $y_{n+1}$ [B batch, T positions, 768 d]
parameters: 9,438,720 · × 3 blocks = 28,316,160

Parameters

9,438,720 parameters per block, across its learned components

Exact per-instance accounting
Component	Calculation	Parameters
LayerNorm	$768$	768
Self-attention	2,359,296 across its components	2,359,296
LayerNorm	$768$	768
Gated MLP	7,077,888 across its components	7,077,888
Per-instance total		9,438,720

Across 3 independent blocks: $9{,}438{,}720 \times 3 = 28{,}316{,}160$ 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.

Explanation

It is the denoiser’s block class without the time: no modulation and no gate, the gated MLP in place of the standard one, and no softcap.

In training a dropout of 0.1 acts on each update before it is added.

Sources5

Direct evidence

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