Add noise

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

Add noise

Encoder gradient scale

×

no parameters

Leave the latents’ values as they are and multiply the gradient that returns through them to the encoder by γ.

Why

Two pulls. At the start of joint training the two objectives pull the latent space in different directions: the reconstruction loss encourages the encoder to enlarge the latents so that decoding is easy, and the diffusion loss encourages them to shrink so that denoising is trivial. With both acting from the first step the encoder struggles to learn meaningful representations. LDLM paper §5.2

So the encoder first trains on reconstruction alone and the diffusion objective is brought in gradually. The denoiser itself trains from the beginning; only the gradient of the diffusion loss into the latent encoder is multiplied by γ, which rises from nearly 0 to 1. LDLM paper §5.2

Its form. The scaling mixes the latents with a detached copy of themselves, so in the forward pass the diffusion branch always receives the current latents: the target is unchanged, and only the strength with which the diffusion loss updates the encoder is controlled. The normalized sigmoid keeps γ near its minimum at first and raises it smoothly near the end of the warm-up, so that an early, poorly fitted diffusion model does not strongly affect the encoder before the reconstruction space is stable. LDLM paper Appendix A.2

Its length. Without the warm-up the diffusion objective dominates the encoder’s early updates and quickly reduces the latents’ magnitude. The exact length is not critical once it is long enough to form a decodable latent space; an overly long one delays the diffusion’s shaping of that space. LDLM paper §6

z0γ=γ z0+(1−γ)sg⁡(z0){\htmlData{light=target gradient-scale}{z_0^{\gamma}}}=\gamma\,{\htmlData{light=latents}{z_0}}+(1-\gamma)\operatorname{sg}({\htmlData{light=latents}{z_0}})
in
z0z_0 [B batch, T positions, 768 d]
out
z0γz_0^{\gamma} [B batch, T positions, 768 d]
Explanation

sg⁡\operatorname{sg} stops the gradient. In value z0γ=z0z_0^{\gamma}=z_0; backwards, the diffusion loss reaches the encoder multiplied by γ\gamma.

γ\gamma rises from 0.001 to 1 over the first 50,000 steps by a sigmoid in the step, centred at 80% of the warm-up with steepness 10 and normalized to run from 0 to 1; after that it is 1.

Both the noised latents and the loss’s target are made from this copy, so the scale applies to both ways the loss reaches the encoder.

Sources3

Direct evidence

The diffusion gradient into the encoder: the weight rises from its minimum by a sigmoid schedule during warm-up, and the clean latents become w · z + (1 − w) · detach(z), unchanged in value.

The implementation, pinned at revision e18c42acdfb7, not yet public

Direct evidence

The warm-up schedule: a sigmoid centred at 0.8 of the warm-up, normalized to run from 0 to 1.

The implementation, pinned at revision e18c42acdfb7, not yet public

Direct evidence

Joint training: decoder noise of standard deviation 3; weights 1 for the three losses; the diffusion gradient into the encoder rises from 0.001 over 50,000 steps by a straight-through scaling; the token decoder’s input is detached; AdamW at 3·10⁻⁴ for 1,000,000 steps of 512 sequences, gradients clipped at norm 1.

The implementation, pinned at revision e18c42acdfb7, not yet public