Masked diffusion post-training stops sampling blocks at random: a full-block estimator scores highest on four of five benchmarks

In post-training for masked diffusion language models, the Monte Carlo estimate of the blockwise surrogate objective now covers all blocks, evaluating each exactly once and averaging, instead of drawing the block index at random. The existing estimation scheme randomly samples the block index, so it may evaluate the same block more than once and leave some blocks entirely unevaluated, adding variance from block selection to the estimate. In a simulation with C=4 blocks and K=4 samples, at least one block went unselected in each of 8 independent trials; the derivation gives a probability of about 0.91 that at least one block is never selected when K=C=4.

The change replaces only the estimator while the blockwise surrogate objective stays fixed, so it can be integrated into existing masked diffusion post-training pipelines. Prior work designed policy optimization rules, improved the surrogate objective, or controlled the bias and variance of likelihood approximation; the estimator design itself had not been examined separately. Experiments fine-tune LLaDA-8B-Instruct with QLoRA for 4,000 steps on five benchmarks, GSM8K, MATH-500, Countdown, Sudoku and HumanEval, with all compared methods sharing one hyperparameter setting.

Across the five benchmarks, the full-block estimator reaches the highest accuracy on four and averages 12.79 points above the random block sampling baseline. The full-block estimator beats the random block sampling baseline on all five benchmarks, and its simple average over the five is 5.05 points above the strongest baseline, SPG. The scores are 82.26 on GSM8K, 38.00 on MATH-500, 91.39 on Sudoku and 36.59 on HumanEval; on Countdown it reaches 66.41, 22.27 points above random block sampling but below SPG's 70.31.

The estimation-quality gain is directly measurable at the estimator level: on GSM8K the full-block estimator at K=4 has rank correlation 0.878 and RMSE 50.1 against the reference value. Random block sampling at the same K=4 has rank correlation 0.602 and RMSE 136.0, so the full-block estimator's accuracy is roughly comparable to random block sampling at about 32 samples. The variance ratio ranges from 0.08 to 0.12 across the five datasets, putting the full-block estimator's variance at about one tenth of random block sampling; at the policy-gradient level, however, the ratio ranges from 0.79 to 0.96.

Experiments fine-tune LLaDA-8B-Instruct with QLoRA for 4,000 steps and cover five benchmarks across mathematics, planning and coding, with C=4 blocks and an additional check on GSM8K and MATH-500 at C=8. Accuracy is evaluated only at each method's final checkpoint, with no comparison at intermediate checkpoints.

The authors note that the estimator-level variance reduction does not carry over equally to the policy gradient, whose variance ratio ranges from 0.79 to 0.96, so the accuracy gain is not necessarily due to reduced gradient variance; the full-block estimator also delivers a learning signal to every block at every step, and how much each of the two factors contributes is not disentangled. They also state that computational limits prevented experiments with more training steps or a larger GRPO group size, that combinations with other policy optimization methods such as PPO and DPO were not verified, and that effectiveness on open-ended natural language generation and summarization is unconfirmed; parallel full-block evaluation scales GPU memory linearly with the number of blocks, and extending the idea to continuous or high-dimensional random variables requires approximation.