Sampling reorders three diffusion LMs by bias; scores move over 40%

Beyond Model Weights: Sampling-Dependent Fairness in Diffusion Language Models

Koshiro Saito, Daisuke Oba, Danushka Bollegala, Naoaki Okazaki

cs.CL

2026-10-20

Across 24 setups, LLaDA, LLaDA-1.5, and Dream shift Open-BBQ bias by over 40% and swap rank. SFT cuts LLaDA-8B from 0.123 to 0.015; sampling sensitivity remains.

What problem this solves

Diffusion language models now match autoregressive models on generation quality, and deployment adds three controls autoregressive decoding does not have: block length, remasking, and tokens per step (TPS). Those knobs fix the speed-quality point. Fairness has barely been measured. Temperature and top-p already move bias scores for autoregressive models. Diffusion controls change the generation process itself: which positions are written in parallel, and which predictions get rewritten. A fairness number from one configuration describes the sampler about as much as the weights.

Method

The models are LLaDA-8B-Instruct, LLaDA-1.5, and Dream-7B. Open-BBQ is used in short-answer form: 10,000 items drawn from 58,492, half ambiguous and half disambiguated, with the original mix across 11 social categories. The model writes a rationale, then puts the option index in a box. An LLM judge in the pilot hit its own guardrails, and the failures grew with faster commitment, so the bias score counts only answers whose box parses.

On ambiguous items the correct choice is 'can't be determined'. On disambiguated items the context names the answer. Bias is the gap between stereotype-aligned and opposite answers, divided by the answers that count. Zero means the two sides are chosen equally often. One means the stereotype every time. Correct abstentions shrink the ambiguous score.

The grid has 24 cells on a fixed 256-token window: block length 8, 32, or 128, and TPS 1, 2, 4, or 8. Tokens committed in the same step are predicted together and cannot condition on each other. LLaDA compares low-confidence remasking with random remasking at temperature 0. Dream compares low-entropy remasking with its original sampler at temperature 0.1, because temperature 0 returned empty answers. Dream's official sampler has no blocks; block sampling was added for this grid.

SFT supervises only the 5-token boxed answer, fully masked, so the 1/t factor in the diffusion loss does not blow up the gradient on a short target. dOPSD treats later states on the same trajectory as privileged information, freezes the teacher at the initial policy, and matches it with forward KL. Every rollout is kept. Filtering on correctness would drop the biased answers the update is meant to fix. A fairness penalty, applied only on ambiguous items, pulls log-probabilities of group terms together: a pairwise log-ratio, a JS divergence to the uniform distribution, and within-group variance, with weights 1, 30, and 1. The last two also use a 159-group lexicon over the 11 categories. Training runs for 3 epochs at learning rate 4e-5 with LoRA. On eight H200s, SFT takes about 5 minutes and dOPSD about 951 minutes.

Results

All three models prefer stereotype-aligned answers on ambiguous items, and disambiguated bias sits near zero. For base LLaDA-8B with low-confidence remasking, the 12 block-length and TPS cells fall between -0.018 and -0.005. There is no head-to-head score against an autoregressive model on this split.

Bias scores move by more than 40% across the 24 configurations for every model, and which model looks most biased changes with the cell. Remasking moves LLaDA the most. Higher TPS widens the spread across block lengths and raises the rate of unparsable boxes. Dream-7B's box failures get especially high at large TPS, so those bias numbers are weak. Under greedy low-confidence remasking the trajectory is deterministic, and the score still moves with block length and TPS.

MethodMetricResult
Base, one settingAmbiguous bias0.123
SFTAmbiguous bias0.015
SFT+JS (weight 30)Ambiguous bias0.000
dOPSDAmbiguous bias0.064
dOPSD+JS (weight 30)Ambiguous bias0.011
Base to SFT+JSAmbiguous accuracy0.703 to 0.999
11 categories, unweighted, default samplingMean biasBase 0.171, SFT 0.030, SFT+JS 0.001

Ambiguous accuracy rises to 0.965 under SFT and 0.999 with the JS term because ambiguous items are now answered 'can't be determined'. SFT's disambiguated accuracy rises to 0.965 from a base of 0.942. dOPSD's falls to 0.934, a significant drop. Unweighted across 11 categories, dOPSD sits at 0.108, and at 0.015 with the JS term.

SFT reaches the lower score on most cells. dOPSD moves less across cells. After SFT, LLaDA's rank among the three models runs from first to third. After dOPSD the rank stays mostly put. In Figure 4 the within-remasking range of ambiguous bias on LLaDA-8B is about 0.035 and 0.037 for the base model, 0.154 and 0.163 after SFT, and 0.207 and 0.205 with the log-ratio term. Lexicon penalties bring the dOPSD ranges down to about 0.019-0.033. A single-setting 0.000 does not describe the grid: SFT+JS still has ranges of 0.040 and 0.079. On LLaDA-1.5, SFT plus the variance penalty stretches the range to 0.222 and 0.231. Cutting the average and cutting sampling sensitivity are separate jobs.

Utility is measured only at the default sampler: block length 32, TPS 1, low-confidence remasking. GSM8K uses block length 8. Bias falls significantly. IFEval and MMLU show no significant loss, and SFT leaves GSM8K intact. LLaDA-1.5 and Dream-7B each have at least one recipe with no significant loss on any of the three benchmarks. dOPSD hurts GSM8K on both, and on Dream-7B it also hurts IFEval, with two variants hurting MMLU. Base bias concentrates on age, physical appearance, socioeconomic status, and disability. Category rank stays aligned with the base model, Spearman correlation at least 0.72, aside from two near-zero columns.

Why it matters

Anyone shipping a diffusion LM is already choosing these three knobs for speed. A fairness report from the default cell can change, in level and in model order, when TPS changes.

Penalties on the per-position distribution transfer, because every masked position has one. Inference-time interventions, parameter edits, and preference tuning need a sequence likelihood or a next-token distribution, which these models do not expose directly.

SFT on eight H200s takes about five minutes and moves default ambiguous bias from 0.123 to 0.015, with no significant drop on IFEval or MMLU. The training recipes are familiar. The reporting rule is what changes: after debiasing, the speed settings you intend to ship still have to be swept.

Limitations

The study covers three models, 11 Open-BBQ categories, and a finite sampler grid. The paper states that the size of the sampling effect need not match on other families, benchmarks, or definitions of fairness.

Unparsable answers are dropped before the bias score. On Dream-7B at high TPS, an extreme score often means the box never appeared. Forcing a boxed index also separates this protocol from Open-BBQ's free-form answers plus an LLM judge. Utility was not re-run across the 24 cells, so a claim of no downstream cost holds at one operating point.

SFT supervises little besides the option index, so the model may mostly have learned which digit to box on this item set. The lexicon groups healthy with old and young, and plumber with gay and queer. Equalization was measured on that lexicon only. Each model was trained once.

Terms

Source

What people are saying

Related papers

All paper explainers