Summary: Train the Sampler You Run — Few-Step Masked Diffusion Language Models via Mixture-Likelihood Fine-Tuning

Author: Amin Najafgholizadeh (University of Zanjan) · arXiv:2611.02666v1 · Preprint, November 3, 2026

---

Summary (Overview)

• Core contribution: Introduces mixture-likelihood fine-tuning (Mix-FT), a teacher-free, architecture-preserving fine-tune that trains a pretrained masked/block diffusion language model on an importance-weighted (IWAE) bound of its own few-step sampler's likelihood, rather than on per-token cross-entropy.
• Key insight: A $K$-step random-order sampler draws from a mixture over reveal plans $q_\theta(x) = \sum_d \mu(d)\, q_{\theta,d}(x)$. Cross-entropy (CE) averages $\log q_{\theta,d}$ over plans and is minimized by the true conditionals — which do not minimize $-\mathbb{E}_p \log q_\theta$. Mix-FT instead optimizes $-\log$ of the plan average.
• Efficiency: On LLaDA-8B, a LoRA fine-tune costing 0.002% of pretraining compute beats continued CE and Alpha diffusion at every tested step count. With half the model evaluations (16 vs. 32 steps per block), its continuations win under three judges, and two blinded LLM raters prefer them in 69% and 71% of pairs.
• Public 110M model (BD3-LM): Under four T4-hours, Mix-FT beats continued CE, Alpha diffusion, and SDTT distillation at 4, 8, and 16 steps per block (0.25–0.45 nats/token over the strongest baseline). With $8\times$ budget, its 8-step samples close 95% of the original model's 8→16-step gap in ~half the generation time.
• Mechanism (exact text8 experiments): Using exact $K$-step likelihoods via a dynamic program over reveal sets, about three quarters of the likelihood gain over continued CE is attributed to weighting each plan by how well it explains the sequence (isolated by a uniform-weight control).

---

Introduction and Theoretical Foundation

Background and Motivation

Masked diffusion language models (Sahoo et al., 2024; Shi et al., 2024; Lou et al., 2024; Nie et al., 2025b) fill masked positions conditioned on revealed ones. Their appeal is few-step parallel sampling: a block of $B$ positions is filled in $K < B$ steps, each revealing $B/K$ positions drawn independently given what is already revealed.

The resulting factorization error is well known. Existing fixes act at:

• Sampling time — temperature (Amin et al., 2026), better schedules (Luxembourg et al., 2026; Ben-Hamu et al., 2025);
• Architecture — added heads or latents (Kim & Sung, 2026; Xie et al., 2026; Ivanov et al., 2026);
• Distillation — from a many-step teacher (Deschenaux & Gulcehre, 2024; Hayakawa et al., 2025; Zhang et al., 2026).

This paper changes the loss instead.

Theoretical Foundation

Setup (Section 2): Block diffusion (Arriola et al., 2025) splits a sequence into blocks of $B$ positions, generated left to right. Within a block, a plan $d = (G_1, \dots, G_K)$ is an ordered partition into $K$ groups of size $g = B/K$, drawn uniformly. The denoiser is $r_\theta(x_i \mid x_S)$. The sampler's law is:

$$q_{\theta,d}(x) = \prod_{k=1}^{K} \prod_{i \in G_k} r_\theta\!\left(x_i \mid x_{G_{<k}}\right), \qquad q_\theta(x) = \sum_d \mu(d)\, q_{\theta,d}(x). \tag{1}$$

A temperature $\beta$ replaces $r_\theta$ by $\mathrm{softmax}(\beta \cdot \text{logits})$.

Why CE is not the sampler's optimum:

Proposition 1 (Fixed or state-determined plans). If the plan is fixed, or the next group is a deterministic function of revealed tokens that does not depend on $\theta$, then $-\mathbb{E}_p \log q_{\theta,d}(x) = \sum_k \mathbb{E}_p\!\left[-\sum_{i \in G_k} \log r_\theta(x_i \mid x_{G_{<k}})\right],$ and every term is minimized by $r_\theta(x_i \mid x_S) = p(x_i \mid x_S)$.

This is the chain rule: for such samplers, nothing done to the denoiser helps beyond learning the conditionals; remaining error is the plan's factorization error. (Confidence-based planners, which read the denoiser's own outputs, are not covered.)

Proposition 2 (Per-path loss = mixture loss + plan information). For every $x$, $\mathbb{E}_{d \sim \mu}\!\left[-\log q_{\theta,d}(x)\right] = -\log q_\theta(x) + \mathrm{KL}\!\left(\mu \,\|\, \pi_\theta(\cdot \mid x)\right),$ where $\pi_\theta(d \mid x) = \mu(d)\, q_{\theta,d}(x) / q_\theta(x)$ is the posterior over plans.

The left side (averaged over $p$) is the per-path loss — CE at the sampler's states, up to per-position weighting. It charges the sampler for any information a sample carries about its own plan. No reweighting of masks, steps, or orders removes this term, since every such loss keeps the logarithm inside the plan average.

Two effects explain how the mixture optimum differs from what CE learns:

1. Within-step dependence ($g > 1$): the mixture can let plans specialize (e.g., commit early to a mode so later same-step tokens are nearly determined). Re-optimizing a free table of conditionals for the mixture removes nearly all error left by true conditionals, more than a tuned per-step temperature (Table 23).
2. Fit under the order mixture at one token per step ($K = B$): with unlimited capacity, exact conditionals make every order exact, so CE would be optimal. A finite network cannot make every order agree. The per-path loss (any-order NELBO) asks it to fit every order equally; the mixture likelihood needs only the average over orders to be good, and weights move effort toward orders that explain each sequence.

---

Methodology

Objective: Importance-Weighted Bound

Mix-FT fine-tunes on the random-order sampler's mixture likelihood via the IWAE bound (Burda et al., 2016) over $M$ plans drawn from $\mu$:

$$\mathcal{L}_M(\theta) = \mathbb{E}_{x,\, d_{1:M}}\!\left[-\log \frac{1}{M} \sum_{m=1}^{M} q_{\theta,d_m}(x)\right], \tag{2}$$

an upper bound on $-\mathbb{E}_p \log q_\theta(x)$ that tightens monotonically in $M$.

Cost: $1 + M(K-1)$ forward passes per sequence — one with the block fully masked (shared by all plans) plus $K-1$ per plan. Blocks share passes since they are conditionally independent given the clean prefix.

Gradient (with self-normalized weights $\tilde{w}_m \propto q_{\theta,d_m}(x)$):

$$\nabla_\theta \mathcal{L}_M = -\mathbb{E}\!\left[\sum_m \tilde{w}_m \sum_k \sum_{i \in G_k^{(m)}} \nabla_\theta \log r_\theta\!\left(x_i \mid x_{G_{<k}^{(m)}}\right)\right],$$

i.e., cross-entropy at the sampler's own states, with each plan weighted by how well it explains $x$. At $M = 1$ the weight is 1 and $\mathcal{L}_1$ is the per-path loss (optimum: true conditionals) — so the method requires $M > 1$. As $M \to \infty$, weights approach the plan posterior and the gradient approaches that of $-\log q_\theta(x)$.

In exact toy experiments (Table 24, Section F): training a free table on $\mathcal{L}_1$ leaves it where it started, while $M = 16$ removes 94% of the error of the true conditionals at $K = 3$.

Exact Evaluation via Dynamic Programming

Because $r_\theta$ depends on the revealed set, not the reveal order, the sum over plans in Eq. (1) collapses to a dynamic program over reveal sets (Section G). At $B = 8$, $K = 4$: 127 forward passes per sequence for 2,520 plans. Used for all block-8 text8 results, since IWAE-64 is looser for Mix-FT than for CE and would understate the gain.

Experimental Setup

• 110M (Section 4): Public BD3-LM OpenWebText checkpoint (block 16, 1,024 GPT-2 tokens). All arms start from the same weights with equal extra compute (<0.02% of pretraining, <4 T4-hours). Mix-FT trains on two-step plans ($K=2$) with $M=8$ plans per sequence. Baselines: continued CE at four-step states; Alpha diffusion ($\kappa = 0.5$); SDTT-4 and SDTT-8 distillation. Three seeds per arm.
• Evaluation: 512 samples of 1,024 tokens at six temperatures; scored by GPT-2-large and Pythia-1.4B on content tokens (excluding newline/EOT) at matched per-sample entropy (four levels from real text's value down 0.3 nats); plus MAUVE, Vendi score, self-BLEU-4, and blinded pairwise LLM raters (ChatGPT, Claude Sonnet).
• Scale (Section 5): LLaDA-8B-Base with LoRA (0.0018% of pretraining; block 32, 192-token continuations); 1.1B SMDM base model.
• Text8 (Section 6): 4.82M-parameter block diffusion transformer, block 8, $K=4$; base = masked-CE pretraining + four units of CE at sampler states ("Base 4×"); exact $K$-step NLL in bpc on 2,000 held-out sequences.

---

Empirical Validation / Results

110M Public Model (Table 1)

Judge NLL on content tokens at matched per-sample entropy (nats/token, lower is better); means of 3 seeds.

Key findings:

• Mix-FT gives the best judged samples from 4 to 16 steps, beating the strongest fine-tuned baseline by 0.25–0.45 nats/token and continued CE by 0.52–0.72.
• Every fine-tuned baseline ends behind the untouched start at $K=16$ (continued CE by 0.15–0.25 at $K=8$ and 16); even BD3-LM's own continued objective ends 0.18–0.43 behind the start.
• Two-step training plans beat finer ones: training on $K=2$ beats the same recipe on $K=4$ by 0.12–0.31 nats/token at every step count — and not because of extra updates (a 443-step two-step run still wins; four-step runs with matched updates stay 0.25–0.48 behind at $K=4$). Even training on the $K=8$ plans the sampler runs is worse at $K=8$ than two-step training. Hypothesis: coarse plans give the objective the most to correct — the start's plan weights are most concentrated at $K=2$ (median effective sample size 1.32 vs. 1.54 at $K=4$, 1.95 at $K=8$).
• Likelihood price: NELBO perplexity rises (23.0 → 24.5), but the tighter IWAE-64 bound of the one-token-per-step sampler ends 0.002–0.003 nats/token below continued CE's — exactly the looser per-path fit that effect (ii) predicts.
• Half the steps: With $M=256$ at $8\times$ budget, Mix-FT at $K=8$ closes 95% (GPT-2-large; bootstrap interval 92–99%) and 97% (Pythia) of the start's own 8→16 gap in ~half the generation time. At $16\times$ ($M=512$), its 8 steps beat the start's 16 in every seed with the first token drawn.
• Blinded raters (Fig. 3b): The start's $K=8$ text wins only 6–7% of pairs against its own $K=16$; Mix-FT at $K=8$ wins 45% pooled — recovering 88% of the lost win rate (95% interval 79–97%). At equal steps, raters prefer Mix-FT's text in 76% and 79% of pairs.
• MAUVE: 0.86–0.89 for Mix-FT at $K=8$, vs. 0.68 (Alpha) and 0.55 (start).
• Compute alone does not close the gap: baselines given the full $8\times$ budget stay >0.7 nats/token behind the start at $K=16$, seeing 65–129× as many distinct sequences as $8\times$ Mix-FT.
• Confidence planner transfer: Mix-FT never trained on LLaDA's low-confidence remasking but still beats every baseline and the start at $K=4, 8, 16$ under that planner.

LLaDA-8B (Table 2)

Start at $K=32$ minus Mix-FT at $K=16$ (positive: Mix-FT better):

Mix-FT at $K=16$ (96 forward passes) beats LLaDA's $K=32$ (192 passes) in every seed under all three judges with either planner. MAUVE rises from 0.63 to 0.73–0.77 (confidence decoding) and 0.85 to 0.95 (random order).

Text8: Why It Works (Table 3)

Block 8, $K=4$: exact $K$-step NLL (bpc, 3-seed mean); $\Delta$ vs. continued CE (negative = better); $G$ = judge-NLL gap at matched distinct-4 (positive = better); $\Delta_8$ = exact NLL at $K=8$; $\Delta_{16}$ = block 16, $K=4$, IWAE-64.

Findings:

• A second stage, not a replacement: used from scratch, the mixture objective is worse by ~0.06 bpc after four units, with 25× fewer optimizer steps and noisier gradients. CE learns conditionals cheaply; the mixture objective then corrects them.
• From a CE start it wins: Mix-FT lowers exact NLL by 0.039–0.043 bpc vs. continued CE in every seed — about three times what continued CE gains from its own fifth unit.
• Plan weighting is the main cause: Per-path FT (identical plans, batches, updates, but uniform $1/8$ weights) beats continued CE but loses to Mix-FT in every seed — the weighting accounts for 73–80% of the likelihood gain. The effective sample size of eight plan weights falls from 4.3 (CE) to 3.6 (Mix-FT), showing plans specialize.
• Splitting the gain: At $K=8$ (one token per step), effect (i) is absent, yet Mix-FT still beats CE trained at $K=8$ states by 0.022–0.028 bpc. The $K=4$ gain of 0.042 bpc splits roughly evenly: 0.025 persisting at one token per step (effect ii) and 0.016 from narrowing the $K=4$–$K=8$ gap (effect i). Widening the network 4.8M → 43M shrinks the first part by 0.010–0.011 bpc and leaves the second level — as a capacity account of effect (ii) predicts.

---

Theoretical and Practical Implications

Theoretical:

• Proposition 2 (a special case of Theorem 1 of Amin et al., 2026) cleanly separates the sampler's mixture likelihood from the per-path loss, showing the per-path loss pays an irreducible plan-information KL cost. This reframes CE-based training of parallel samplers as optimizing the wrong objective.
• The two-effect decomposition (within-step dependence; finite-capacity fit under the order mixture) provides a testable account, validated by exact likelihoods, a uniform-weight control, and a capacity-scaling experiment.
• The results challenge the assumption that distillation from many-step teachers is necessary for good few-step generation.

Practical:

• Mix-FT is a drop-in change of loss: no teacher, no architecture change, no sampler change, compatible with LoRA, and costing under 0.02% of pretraining compute.
• It halves inference cost at near-equal quality (LLaDA-8B: 16 steps beat 32; 110M: 8 steps close 95% of the 8→16 gap), directly reducing deployment cost of diffusion LMs.
• Counterintuitive but actionable: train on coarser plans than you sample with ($K=2$ training improves $K=4$–16 sampling).
• Transfers to unseen step counts and to confidence-based decoding without retraining.

---

Conclusion

Few-step masked diffusion samplers are mixtures over reveal plans, and fine-tuning on that mixture's likelihood is a cheap, teacher-free change of loss. Main takeaways:

1. Mix-FT optimizes an importance-weighted bound on the deployed sampler's likelihood; its gradient is CE at the sampler's own states with plans weighted by how well they explain the sequence.
2. On LLaDA-8B, 16 steps per block beat the original's 32 (69%/71% blinded rater preference); on a public 110M model it gives the best samples from 4–16 steps, and with more plans its 8 steps close 95% of the gap to the original's 16 in about half the time.
3. Exact text8 experiments attribute ~3/4 of the likelihood gain to plan weighting, with the remainder split between within-step dependence and finite-capacity fit under the order mixture.

Limitations and open directions:

• Scale: full baselines only at 110M; at 8B, Mix-FT at $K=8$ stays behind LLaDA at $K=16$; on a 1.1B GSM8K self-training setup, Mix-FT did not differ from continued CE (as the fixed point of mixture refit predicts for self-generated data).
• Two-step regime: SDTT-4 still wins at $K=2$, partly by repairing a newline excess Mix-FT keeps.
• Block size: tested 8, 16, 32; with one block over 1,024 tokens, eight plan weights have median effective sample size 1.00, making few-plan bounds loose.
• Price: NELBO perplexity rises (23.0 → 24.5 at 1×, 28.8 at 8×); the method needs a CE-trained start and costs $1 + M(K-1)$ passes per sequence.
• Scope: uniform random-order samplers only (Proposition 1 does not cover confidence planners); PUMBA, Di4C, and added-head methods were not run; judges and raters are models.

"Few-step masked diffusion samplers are mixtures over reveal plans, and fine-tuning on that mixture's likelihood is a cheap, teacher-free change of loss."