Sampler

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

Latents at time tz[B, T, 768]Previous estimateẑsc[B, T, 768]Adaptive time grid1,024 stepst[B]z0[B, T, 768]
ϵ^\hat\epsilon==(zt−αtz^0)/σt({\htmlData{light=state}{z_t}}-\alpha_t{\htmlData{light=prediction denoiser}{\hat z_0}})/\sigma_tss==−ϵ^/σt-\hat\epsilon/\sigma_tff==−12βtzt−βts-\tfrac12\beta_t{\htmlData{light=state}{z_t}}-\beta_tsmt′{\htmlData{light=step-mean mean}{m_{t'}}}==zt+f (t′−t){\htmlData{light=state}{z_t}}+f\,(t'-{\htmlData{light=time grid}{t}})zt′{\htmlData{light=next-state diffuse}{z_{t'}}}==mt′+βt (t−t′)  ξ{\htmlData{light=step-mean mean}{m_{t'}}}+\sqrt{\beta_t\,({\htmlData{light=time grid}{t}}-t')}\;\xi
mt′[B, T, 768]zt′[B, T, 768]starts as Gaussian noisereplaced by the next latentsstarts as zerosreplaced by this call’s predictionthe next call’s latentsreturned by the last callcall 1 of 1,024 · m: the step’s mean, before its noise · t′: the grid’s next time · α, σ, β: the schedule at t · ξ: fresh Gaussian noise

Sampler

learned

12 transformer blocks that predict the clean latents from the noised ones, the diffusion time and, when there is one, their own previous estimate.

xout=xˉ⊙(1+so)+box_{\mathrm{out}}=\bar x\odot(1+s_o)+b_o
z^0=xoutWout+bout\hat z_0=x_{\mathrm{out}}W_{\mathrm{out}}+b_{\mathrm{out}}
All 30 steps
x0=[ zt  z^0sc ]Win+binx_0=[\,z_t\;\hat z_0^{\mathrm{sc}}\,]W_{\mathrm{in}}+b_{\mathrm{in}}
Time embedding:\text{Time embedding:}
et=[cos⁡(t ωj)]j ∥ [sin⁡(t ωj)]je_t=\big[\cos(t\,\omega_j)\big]_{j}\,\Vert\,\big[\sin(t\,\omega_j)\big]_{j}
ωj=10−j/384,  j=0,…,383\omega_j=10^{-j/384},\;j=0,\ldots,383
e↑=etW1+b1e^{\uparrow}=e_tW_{1}+b_{1}
eσ=e↑ σ(e↑)e^{\sigma}=e^{\uparrow}\,\sigma(e^{\uparrow})
e↓=eσW2+b2e^{\downarrow}=e^{\sigma}W_{2}+b_{2}
at=e↓ σ(e↓)a_t=e^{\downarrow}\,\sigma(e^{\downarrow})
Denoiser block:\text{Denoiser block:}
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}
[ 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}
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}
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}}
xˉ=γout⊙x12−μ(x12)σ2(x12)+ε\bar x=\gamma_{\mathrm{out}}\odot\frac{x_{12}-\mu(x_{12})}{\sqrt{\sigma^2(x_{12})+\varepsilon}}
ε=10−5\varepsilon=10^{-5}
mout=[ bo  so ]=atWo+com_{\mathrm{out}}=[\,b_o\;s_o\,]=a_tW_{o}+c_{o}
xout=xˉ⊙(1+so)+box_{\mathrm{out}}=\bar x\odot(1+s_o)+b_o
z^0=xoutWout+bout\hat z_0=x_{\mathrm{out}}W_{\mathrm{out}}+b_{\mathrm{out}}
in
ztz_t [B batch, T positions, 768 d]z^0sc\hat z_0^{\mathrm{sc}} [B batch, T positions, 768 d]tt [B batch]
out
z^0\hat z_0 [B batch, T positions, 768 d]
parameters
132,836,352

Generation runs the averaged weights in evaluation mode: no dropout. Only the stack is under bfloat16 autocast here; the time embedding and the input projection run in float32.

Parameters

132,836,352 parameters per instance, across its learned components

Exact per-instance accounting
ComponentCalculationParameters
Input projection1,536×768+7681{,}536 \times 768 + 7681,180,416
Time embedding2,361,600 across its components2,361,600
Denoiser block10,626,816×1210{,}626{,}816 \times 12127,521,792
Final LayerNorm768768768
Time projection768×1,536+1,536768 \times 1{,}536 + 1{,}5361,181,184
Output projection768×768+768768 \times 768 + 768590,592
Per-instance total132,836,352

Across 1024 shared calls: 132,836,352×1=132,836,352132{,}836{,}352 \times 1 = 132{,}836{,}352 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 predicts the clean latents themselves; the noise estimate and the score are computed from that prediction where the sampler needs them.

The stack runs under bfloat16 autocast; in training the whole call does.

Generation uses an exponential moving average of these weights, kept by the torch_ema library with a configured decay of 0.9999. How that library averages is not recorded here.

Sources4

Direct evidence

The denoiser owns the time embedding’s two affine layers, one rotary table for all blocks, an affine projection of the concatenated input, and the stack.

The implementation, pinned at revision e18c42acdfb7, not yet public

Direct evidence

The time embedding passes its MLP and a SiLU; the noised latents are concatenated with the previous clean estimate, zeros when there is none, and projected; the stack runs in bfloat16.

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