Denoiser

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

Denoiser

Denoiser block

Open inside →×

learned

One of 12 blocks: self-attention and an MLP on the latents, each normalized, then scaled and shifted by the time, and each update gated by the time before it is added.

Δ~mlp=gm⊙Δmlp\tilde\Delta_{\mathrm{mlp}}=g_{m}\odot \Delta_{\mathrm{mlp}}
xn+1=xn′+Δ~mlpx_{n+1}=x'_n+\tilde\Delta_{\mathrm{mlp}}
All 24 steps
mn=[ ba  sa  ga  bm  sm  gm ]=atWn+cnm_n=[\,b_a\;s_a\;g_a\;b_m\;s_m\;g_m\,]=a_tW_{n}+c_{n}
xˉattn=γ1⊙xn−μ(xn)σ2(xn)+ε\bar x_{\mathrm{attn}}=\gamma_{1}\odot\frac{x_n-\mu(x_n)}{\sqrt{\sigma^2(x_n)+\varepsilon}}
ε=10−5\varepsilon=10^{-5}
xattn=xˉattn⊙(1+sa)+bax_{\mathrm{attn}}=\bar x_{\mathrm{attn}}\odot(1+s_{a})+b_{a}
Self-attention:\text{Self-attention:}
[ q  k  v ]=xattn WQKV[\,{\htmlClass{vl-stream-q}{q}}\;{\htmlClass{vl-stream-k}{k}}\;{\htmlClass{vl-stream-v}{v}}\,]=x_{\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}
S=QKTS=QK^{\mathsf T}
F=50tanh⁡ ⁣(S5064)F=50\tanh\!\Big(\frac{S}{50\sqrt{64}}\Big)
At,s=exp⁡Ft,s∑u=1Texp⁡Ft,uA_{t,s}=\frac{\exp F_{t,s}}{\sum_{u=1}^{T}\exp F_{t,u}}
O=AVO=AV
o=reshape⁡(O)o=\operatorname{reshape}(O)
Δattn=oWO\Delta_{\mathrm{attn}}=oW_{O}
Δ~attn=ga⊙Δattn\tilde\Delta_{\mathrm{attn}}=g_{a}\odot \Delta_{\mathrm{attn}}
xn′=xn+Δ~attnx'_n=x_n+\tilde\Delta_{\mathrm{attn}}
xˉmlp=γ2⊙xn′−μ(xn′)σ2(xn′)+ε\bar x_{\mathrm{mlp}}=\gamma_{2}\odot\frac{x'_n-\mu(x'_n)}{\sqrt{\sigma^2(x'_n)+\varepsilon}}
ε=10−5\varepsilon=10^{-5}
xmlp=xˉmlp⊙(1+sm)+bmx_{\mathrm{mlp}}=\bar x_{\mathrm{mlp}}\odot(1+s_{m})+b_{m}
MLP:\text{MLP:}
h=xmlpW1+b1h=x_{\mathrm{mlp}}W_{1}+b_{1}
g=h2[1+tanh⁡ ⁣(2π(h+0.044715h3))]g=\frac{h}{2}\left[1+\tanh\!\left(\sqrt{\frac{2}{\pi}}\left(h+0.044715h^3\right)\right)\right]
Δmlp=gW2+b2\Delta_{\mathrm{mlp}}=gW_{2}+b_{2}
Δ~mlp=gm⊙Δmlp\tilde\Delta_{\mathrm{mlp}}=g_{m}\odot \Delta_{\mathrm{mlp}}
xn+1=xn′+Δ~mlpx_{n+1}=x'_n+\tilde\Delta_{\mathrm{mlp}}
in
x0x_0 [B batch, T positions, 768 d]ata_t [B batch, 768 features]
out
xn+1x_{n+1} [B batch, T positions, 768 d]
parameters
10,626,816 · × 12 blocks = 127,521,792
Parameters

10,626,816 parameters per block, across its learned components

Exact per-instance accounting
ComponentCalculationParameters
Time projection768×4,608+4,608768 \times 4{,}608 + 4{,}6083,543,552
LayerNorm768768768
Self-attention2,359,296 across its components2,359,296
LayerNorm768768768
MLP4,722,432 across its components4,722,432
Per-instance total10,626,816

Across 12 independent blocks: 10,626,816×12=127,521,79210{,}626{,}816 \times 12 = 127{,}521{,}792 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.

Sources4

Direct evidence

A block owns two gain-only LayerNorms, a fused query–key–value projection and an output projection without bias, an MLP of the configured kind, and, with modulation, an affine layer from the time condition to six vectors, initialized at zero.

The implementation, pinned at revision e18c42acdfb7, not yet public

Direct evidence

A block normalizes, modulates by the time condition’s shift and scale, projects to queries, keys and values, rotates queries and keys, attends over all positions with softcap, and adds the projected output scaled by a gate; the MLP branch does the same with its own shift, scale and gate.

The implementation, pinned at revision e18c42acdfb7, not yet public

Direct evidence

The denoiser’s stack applies the blocks in order and then the output layer, each with the time condition.

The implementation, pinned at revision e18c42acdfb7, not yet public

Direct evidence

Diffusion: self-conditioning on, `t_min` 0.05 and `eps` 10⁻⁵; the adaptive time sampler has 100 bins, a minimum weight of 0.0001, an EMA of 0.9 and updates every 5,000 steps; the denoiser’s EMA decay is 0.9999; scheduler `tanh` with d = 3, solver `euler`, `ode_sampling` false, N = 128 by default; the denoiser has 12 blocks of width 768 with 12 heads, time conditioning `adaln`, a standard MLP, softcap 50 and self-conditioning by concatenation; the timestep embedding’s `max_period` is 10.

The implementation, pinned at revision e18c42acdfb7, not yet public