Stop-gradient
×no parameters
Pass the reconstructed features on unchanged and let no gradient return through them.
Why
Stop. The token decoder’s gradient is stopped at the reconstructed hidden states so that training it does not affect the other networks. LDLM paper §4.1
To speed up training, the token decoder can be left out while the diffusion model trains and tuned after the other networks have converged. The paper states this costs no final quality; its Appendix B.2 reports comparable quality for post-trained decoders of two or more layers and a clearly weaker one of a single layer. LDLM paper §4.1
Size. Joint training uses nine decoder layers, six for hidden states and three for tokens: a large decoder is mainly useful there, where it helps train the latent space through hidden-state reconstruction. LDLM paper Appendix B.2
- in
- [B batch, T positions, 768 d]
- out
- [B batch, T positions, 768 d]
Explanation
So the token loss trains the token decoder alone: neither the latent decoder nor the encoder receives its gradient (detach_decoder_head is true).
Sources2
Direct evidence
Reconstruction: Gaussian noise of fixed standard deviation is added to the latents, the latent decoder runs, the token decoder reads its output detached, and the cross-entropy against the token IDs is taken; the latent decoder’s output is also compared with the standardized features by a mean squared error.
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