Latent encoder

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 encoder

Encoder block

Open inside →×

learned

One of 6 blocks: Perceiver attention from the latents to the latents and the features, then an MLP, each added to the latents.

un′=un+Δattnu'_n=u_n+\Delta_{\mathrm{attn}}
un+1=un′+Δmlpu_{n+1}=u'_n+\Delta_{\mathrm{mlp}}
All 27 steps
Perceiver attention:\text{Perceiver attention:}
uˉn=γa⊙unmax⁡ ⁣(∥un∥2/768,  ε)\bar u_n=\gamma_{a}\odot\frac{u_n}{\max\!\big(\lVert u_n\rVert_2/\sqrt{768},\;\varepsilon\big)}
ε=10−5\varepsilon=10^{-5}
hˉp=γb⊙hpmax⁡ ⁣(∥hp∥2/768,  ε)\bar h^{p}=\gamma_{b}\odot\frac{h^{p}}{\max\!\big(\lVert h^{p}\rVert_2/\sqrt{768},\;\varepsilon\big)}
ε=10−5\varepsilon=10^{-5}
pa=[ q  ka  va ]=uˉn WQKVp^{a}=[\,{\htmlClass{vl-stream-q}{q}}\;{\htmlClass{vl-stream-k}{k^{a}}}\;{\htmlClass{vl-stream-v}{v^{a}}}\,]=\bar u_n\,W_{QKV}
pb=[ kb  vb ]=hˉp WKVp^{b}=[\,{\htmlClass{vl-stream-k}{k^{b}}}\;{\htmlClass{vl-stream-v}{v^{b}}}\,]=\bar h^{p}\,W_{KV}
q=p1:768a{\htmlClass{vl-stream-q}{q}}=p^{a}_{1:768}
Q=γQ⊙Qhmax⁡ ⁣(∥Qh∥2/64,  ε){\htmlClass{vl-stream-q}{Q}}=\gamma_{{\htmlClass{vl-stream-q}{Q}}}\odot\frac{{\htmlClass{vl-stream-q}{Q^{h}}}}{\max\!\big(\lVert {\htmlClass{vl-stream-q}{Q^{h}}}\rVert_2/\sqrt{64},\;\varepsilon\big)}
ε=10−5\varepsilon=10^{-5}
ka=p769:1536a{\htmlClass{vl-stream-k}{k^{a}}}=p^{a}_{769:1536}
kb=p1:768b{\htmlClass{vl-stream-k}{k^{b}}}=p^{b}_{1:768}
k=[ ka;  kb ]{\htmlClass{vl-stream-k}{k}}=[\,{\htmlClass{vl-stream-k}{k^{a}}};\;{\htmlClass{vl-stream-k}{k^{b}}}\,]
K=γK⊙Khmax⁡ ⁣(∥Kh∥2/64,  ε){\htmlClass{vl-stream-k}{K}}=\gamma_{{\htmlClass{vl-stream-k}{K}}}\odot\frac{{\htmlClass{vl-stream-k}{K^{h}}}}{\max\!\big(\lVert {\htmlClass{vl-stream-k}{K^{h}}}\rVert_2/\sqrt{64},\;\varepsilon\big)}
ε=10−5\varepsilon=10^{-5}
va=p1537:2304a{\htmlClass{vl-stream-v}{v^{a}}}=p^{a}_{1537:2304}
vb=p769:1536b{\htmlClass{vl-stream-v}{v^{b}}}=p^{b}_{769:1536}
v=[ va;  vb ]{\htmlClass{vl-stream-v}{v}}=[\,{\htmlClass{vl-stream-v}{v^{a}}};\;{\htmlClass{vl-stream-v}{v^{b}}}\,]
S=QKTS=QK^{\mathsf T}
F=50tanh⁡ ⁣(S5064)F=50\tanh\!\Big(\frac{S}{50\sqrt{64}}\Big)
At,s=exp⁡Ft,s∑u=12Texp⁡Ft,uA_{t,s}=\frac{\exp F_{t,s}}{\sum_{u=1}^{2T}\exp F_{t,u}}
O=AVO=AV
o=reshape⁡(O)o=\operatorname{reshape}(O)
Δattn=oWO\Delta_{\mathrm{attn}}=oW_{O}
un′=un+Δattnu'_n=u_n+\Delta_{\mathrm{attn}}
MLP:\text{MLP:}
uˉn′=LayerNorm⁡(un′)\bar u'_n=\operatorname{LayerNorm}(u'_n)
h=uˉn′ W1+b1h=\bar u'_n\,W_{1}+b_{1}
g=h Φ(h)g=h\,\Phi(h)
Δmlp=g W2+b2\Delta_{\mathrm{mlp}}=g\,W_{2}+b_{2}
un+1=un′+Δmlpu_{n+1}=u'_n+\Delta_{\mathrm{mlp}}
in
u0u_0 [B batch, T positions, 768 d]hph^{p} [B batch, T positions, 768 d]
out
un+1u_{n+1} [B batch, T positions, 768 d]
parameters
8,264,576 · × 6 blocks = 49,587,456
Parameters

8,264,576 parameters per block, across its learned components

Exact per-instance accounting
ComponentCalculationParameters
Perceiver attention3,540,608 across its components3,540,608
MLP4,723,968 across its components4,723,968
Per-instance total8,264,576

Across 6 independent blocks: 8,264,576×6=49,587,4568{,}264{,}576 \times 6 = 49{,}587{,}456 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

Neither branch has a normalization in the block itself: the attention normalizes both of its streams and the MLP its own input.

Sources3

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

The encoder adds positions to the features, normalizes the learned queries, applies the blocks, each attending to the features, and normalizes the result.

The implementation, pinned at revision e18c42acdfb7, not yet public

Direct evidence

The latent encoder: 12 heads of 64 features, query and key normalization, softcap 50, 6 layers, MLP multiplier 4, dropout 0.1, ε = 10⁻⁵.

The implementation, pinned at revision e18c42acdfb7, not yet public