Online Timestep Allocation with Control Variates for Masked Diffusion Language Models

Masked diffusion language models exhibit substantial variation across corruption timesteps, mask patterns, and data samples, making uniform timestep sampling potentially inefficient. We propose Importance Sampling with Control Variates (ISCV), an online, single-forward controller that adapts a full-support timestep proposal using confidence statistics already available in the standard model pass. We show that, with a control variate, the variance-optimal proposal is determined by the second moment of the residualized loss rather than the raw loss. This motivates ISCV-SG, which stop-gradients the correction so that the residual statistics guide future timestep allocation while preserving the expected gradient of the original objective under exact importance weighting. Across five LoRA adaptation tasks, ISCV-SG improves the three-seed mean over uniform training and remains competitive with strong variance-reduction baselines, while retaining runtime close to Standard and requiring approximately half the training time of the strongest baseline P-POTS+MIRROR. Matched ablations and controlled full-parameter pretraining further support the proposed allocation mechanism beyond downstream LoRA adaptation.