# Train the Sampler You Run: Few-Step Masked Diffusion Language Models via Mixture-Likelihood Fine-Tuning > Mix-FT fine-tunes diffusion language models on their own few-step sampler's mixture likelihood, halving inference steps while matching or beating original quality at under 0.002% of pretraining compute. - **Source:** [arXiv](https://arxiv.org/abs/2611.02666) - **Published:** 2026-11-07 - **Permalink:** https://picx.dev/p/vj0MI0 - **Whiteboard:** https://picx.dev/p/vj0MI0/image ## Summary # 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_{ **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_{ 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_{ 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) | model | opt. steps | sampler bound $K{=}4$ | $K{=}2$ | NELBO PPL | judge NLL $K{=}2$ (55 s) | $K{=}4$ (93 s) | $K{=}8$ (162 s) | $K{=}16$ (305 s) | |---|---|---|---|---|---|---|---|---| | Start (BD3-LM) | – | 3.223 | 3.775 | 23.00 | 7.31 | 6.10 | 4.60 | 3.71 | | *first token fixed* † | | | | | 7.46 | 6.17 | 4.54 | 3.63 | | Continued CE | 5,760 | 3.209 | 3.764 | 22.41 | – | 6.20 | 4.82 | 3.89 | | SDTT-4 | 2 × 1,920 | 3.245 | 3.796 | 24.41 | 6.68 | 5.83 | 4.75 | 4.02 | | SDTT-8 | 3,840 | 3.228 | 3.775 | 23.42 | – | – | 4.58 | 3.80 | | Alpha ($\kappa{=}0.5$) | 5,760 | 3.202 | 3.758 | 23.16 | 7.16 | 5.79 | 4.54 | 3.78 | | **Mix-FT (ours)** | **1,152** | **3.181** | **3.723** | 24.51 | **7.07** | **5.54** | **4.10** | **3.37** | | *first token fixed* † | | | | | – | 5.58 | 4.10 | 3.39 | | Real held-out text | | | | | 2.78 | | | | *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) | model | $K{=}8$ random | 16 | 32 | $K{=}8$ low-conf | 16 | 32 | |---|---|---|---|---|---|---| | Start | 5.01 | 3.98 | 3.54 | 5.72 | 5.02 | 4.48 | | Continued CE | 4.88 | 3.89 | – | 5.69 | 4.83 | – | | Alpha ($\kappa{=}0.5$) | 4.67 | 3.74 | – | 5.55 | 4.76 | – | | **Mix-FT (ours)** | **4.29** | **3.48** | – | **5.17** | **4.38** | – | **Start at $K=32$ minus Mix-FT at $K=16$** (positive: Mix-FT better): | judge | random order | low confidence | |---|---|---| | GPT-2-large | +0.07 [−0.02, +0.15] | +0.10 [+0.01, +0.21] | | Pythia-1.4B | +0.10 [+0.02, +0.19] | +0.13 [+0.04, +0.23] | | Qwen2.5-7B | +0.17 [+0.09, +0.25] | +0.13 [+0.04, +0.22] | | ChatGPT rater | – | **69%** [63, 75] | | Claude rater | – | **71%** [66, 77] | 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. | arm | NLL | $\Delta$ | $G$ | $\Delta_8$ | $\Delta_{16}$ | |---|---|---|---|---|---| | Base 4× (start) | 1.804 | | | | | | Continued CE | 1.792 | 0 | 0 | 0 | 0 | | SDTT | 1.802 | +0.009 | −0.038 | – | +0.009 | | Per-path FT | 1.782 | −0.010 | +0.014 | – | – | | Hybrid | 1.764 | −0.028 | +0.076 | – | −0.031 | | **Mix-FT** | **1.751** | **−0.042** | **+0.156** | **−0.025** | **−0.039** | **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." --- _Markdown view of https://picx.dev/p/vj0MI0, served by PicX — AI-generated visual whiteboard summaries of research papers._