XM on text diffusion: worse likelihood, better low-step sampling

We wanted to know whether Explorative Modeling helps a text diffusion model, so we built one and tried it.

The corpus is 579,508 documents, 1.81 billion characters, deduplicated, with 4,000 documents held out for validation and 4,000 for test. The tokenizer is a 32,768-token byte-level BPE trained on the training split. Tokenized, training data comes to 390.3M tokens at 4.57 characters per token.

The model is a 115.7M-parameter masked diffusion language model — 12 layers, 12 heads, d_model 768, SwiGLU feed-forward, RoPE, shared adaLN — trained from random initialisation with continuous-time absorbing-state MDLM (Sahoo et al. 2024) on a linear schedule.

We trained it twice on one RTX 4090: a plain baseline for 20,000 steps in 8h 53m, and XM at K=3 for 10,000 steps in 9h 01m. XM costs 2x per step, so those are the same compute, and that is the pairing.

Both arms were evaluated on test-set likelihood and on generative perplexity across sampling budgets from 8 denoising steps up to 256, scored by a frozen external model. The two measurements point different ways.

1. Likelihood: baseline 3.844 NELBO, XM 4.037.
2. Generation at T=8: XM 15.2% lower generative perplexity, narrowing to 0.3% by T=256.

The rest of this post is the breakdown, one result we did not expect, and what we have not finished checking.

Setup in detail

Two arms, matched on wall-clock rather than on steps:

XM at K=3 costs (K+3)/3 = 2.0 baseline steps: K forward passes under no_grad to score the candidates, then a full forward and backward for the winner, against a baseline step of one forward plus one backward. Measured: 3.25 s/step vs 1.60, a ratio of 2.03x against the predicted 2.00x. So 10,000 XM steps and 20,000 baseline steps are the same spend, and that is the pairing.

Peak VRAM was 17.7 GB for both arms, identical to the digit. Recomputing the winner rather than retaining K graphs (--xm-chunk 1, the reference's save_mem_mode) fully absorbs the exploration cost. XM is free in memory and 2x in time.

One deviation from the reference

The reference NLP branch selects one of K embeddings appended past the vocab and prepends it as an extra sequence position. We instead add it to the timestep conditioning vector that feeds adaLN. The reason is that prepending takes the sequence from 512 to 513 and shifts every RoPE position relative to the baseline, which would confound the comparison; adding to c leaves the architecture byte-identical between arms, so --xm-k is the only difference.

We do not know whether this matters. It is the first thing we would ask the authors about.

We also evaluate all K latents for every sample rather than drawing random candidates per chunk, which makes the selection an exact argmin rather than a K-sample estimate of one. Lower variance, and K is small enough that there is no reason to sample.

Likelihood: XM loses

0.19 nats, same nine hours. Not marginal.

Generation: XM wins, and only where it matters

Generative perplexity under a frozen scorer (gemma-3-270m, revision 9b0cfec8), 1024 samples per setting, 256 tokens each, temperature 1.0, no top-p or top-k, seed 1337.

Six settings, monotone ordering, largest effect where the sampler has the least budget. Noise does not sort itself into a monotone sequence across six independent settings, so we believe the effect is real — but see the pending list, we do not yet have confidence intervals on it.

It is not bought by narrowing. The usual way to fake a generative perplexity win is to collapse toward high-probability text. Both models are in fact too diverse relative to real text, and XM is less so on both axes:

XM moves toward the data on both while also improving perplexity. The entropy gap of 0.098 nats is roughly 3x our measurement floor at 1024 samples (+/- 0.03), so unlike smaller runs we have looked at, this one is actually resolvable.

As a sampling-efficiency claim, it is more modest than 15% sounds. Interpolating the baseline curve to find where it matches XM's quality:

xm_k3 @ T=8   equals base @ T ~ 9.8    1.2x fewer sampling steps
xm_k3 @ T=16  equals base @ T ~ 19.1   1.2x fewer sampling steps

The curve is steep at low T, so a 15% perplexity gain converts to about 1.2x in steps. Inference cost per step is identical between the arms — generation samples one latent and runs one forward — so that 1.2x is a real deployment saving, paid for with 2x training compute.

The part we cannot explain

The intuitive story for XM is that the K latents specialise into distinct modes, and best-of-K training lets each document be explained by whichever mode fits it.

That is not what happened. The latents are redundant.

1. Usage settles at 0.33 / 0.33 / 0.33. No latent dominates, no latent dies.
2. Best-of-3 beats the plain mixture bound by 0.002 nats (4.0354 vs 4.0393).

Read that second number again. We let the model look at the true loss and pick whichever of its three latents scored best on that specific sample, and the oracle peek bought two thousandths of a nat. The three latents converged to nearly identical predictive distributions.

And XM improved low-step sampling anyway.

So whatever the mechanism is, it does not appear to be mode specialisation. Our current guess is that it is a training-dynamics effect rather than a representational one — the argmin acts as a per-sample loss filter, and the gradient the model receives is systematically the easier of three draws. That would be a curriculum effect, not a mixture effect, and it would explain why the benefit shows up in sampling robustness rather than in likelihood. We have not tested this.

If someone has a better hypothesis, we would like to hear it.

Implementation note that will save you a day

wmax — the largest per-latent share of argmin wins over a logging window — reads like total collapse for the first ~50 steps of every XM run, then breaks symmetry, relapses, and breaks again:

step  10   0.935 / 0.029 / 0.036
step  20   0.995 / 0.004 / 0.002
step  30   0.998 / 0.002 / 0.001
step  40   0.997 / 0.003 / 0.000
step  50   0.984 / 0.009 / 0.007
step  60   0.486 / 0.341 / 0.173     <- symmetry breaks here
step  80   0.750 / 0.131 / 0.119     <- and partially relapses
step 110   0.889 / 0.045 / 0.065
step 140   0.443 / 0.299 / 0.258
step 200   0.429 / 0.269 / 0.302
...
step 10000 0.361 / 0.347 / 0.292

The cause is adaLN-zero initialisation. At step 0 the conditioning vector has no effect on the output, so all K latents score identically and argmin returns index 0 for every sample. Losing latents receive no gradient at all in our setup (the latent rows live in reserved vocab slots that are forced to -inf in the logits, so they train only through cond), which looks like a permanent dead-unit trap on paper. It isn't. Once adaLN leaves zero the winner's embedding moves, conditioning starts to matter, and the untouched rows become competitive on a subset of samples.

Practical consequence: do not smoke-test XM for 30 steps and conclude your implementation is broken. The transient runs to roughly step 150 and is not monotone. Read wmax after step 200.

What is still pending

We are publishing this early because the dissociation is worth talking about, not because it is finished. Open, in order of how much they could change the story:

1. A single seed per arm. n=1. There are no confidence intervals in this post. The 15% figure is a point estimate and should be read as one.
2. The run-length confound is not yet excluded. base ran 20,000 steps, xm_k3 ran 10,000. Both completed their own cosine schedule, so both are fully annealed — but we cannot yet rule out that a fully-annealed 10,000-step model simply samples better at low T than a fully-annealed 20,000-step one, for reasons unrelated to XM. A base_10k control sharing xm_k3's exact schedule is running as we write this, and we will append the result. If it beats base at T=8, this post's headline belongs to the learning rate schedule and not to XM.
3. There is already a hint that this confound is live: measured at matched steps, XM's likelihood edge over base shrinks substantially when compared against a 10,000-step baseline instead of a 20,000-step one.
4. One scale, one K. 115M parameters, K=3. No sweep. Whether the effect survives to K=8, or to a model size where the generated text is actually worth reading, is unknown.
5. Absolute quality is bad. Generative perplexity 218 at T=256 against real text's 27. Both arms produce output nobody would want to read. A relative improvement between two incoherent models may not survive into the regime anyone cares about.

Numbers, configs and the exact spec hashes for both runs are reproducible from the run directories; every run logs a controlled_hash covering everything the two arms are required to share, so the comparison is machine-checkable rather than asserted.

More on XM: https://explorative-modeling.github.io/

Appending base_10k when it completes training.