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

Feature reconstruction

×

no parameters

How far the reconstructed features are from the standardized GPT-2 features.

Why

Why MSE. With noise added to the decoder’s input, a cross-entropy loss forces all latents to be well separated, so that the decoder can name the tokens exactly. The mean squared error on hidden states is less strict: two close latents probably have close hidden states, and an average of the two already gives a low loss. It does not strongly force the latents apart, and it still prevents the latent space from collapsing. LDLM paper §5.1

It is the only reconstruction signal that reaches the latent encoder. LDLM paper §4.1

In the ablation, run with the token decoder’s stop-gradient switched off, the hidden-state loss gives substantially better generation than token cross-entropy, and adding cross-entropy to it makes generation worse: the paper reads this as token-level supervision conflicting with the smooth latent geometry diffusion needs. LDLM paper §6

Lh=1∣A∣∑(b,t)∈A1768∥h^b,t−hb,t∥22\mathcal L_{h}=\frac{1}{\lvert A\rvert}\sum_{(b,t)\in A}\frac{1}{768}\lVert\hat h_{b,t}-h_{b,t}\rVert_2^2
in
h^\hat h [B batch, T positions, 768 d]hh [B batch, T positions, 768 d]
out
Lh\mathcal L_{h} [1 one value]
Explanation

The target is detached, and GPT-2 is frozen in any case.

This is the only reconstruction signal that reaches the encoder: it passes back through the latent decoder and the decoder noise.

Sources4

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

Cross-entropy and mean squared error are averaged over valid positions.

The implementation, pinned at revision e18c42acdfb7, not yet public

Direct evidence

With the feature model frozen, its hidden state is computed without gradient, standardized, and passed to the encoder; the same tensor is the target of the feature loss.

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