Adaptive Order Policies for Masked Diffusion

• URL: http://arxiv.org/abs/2606.00295v1
• Date: Fri, 29 May 2026 19:26:53 GMT
• Title: Adaptive Order Policies for Masked Diffusion
• Abstract summary: We propose a scheme for learning the unmasking order using an additional lightweight policy network on top of a diffusion model.<n>Our proposed loss reweights terms in the masked diffusion loss according to policy probabilities, and results in a policy that prefers positions where the denoiser is more likely to be correct.
• Score: 56.27207032553083
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Masked diffusion models have seen great success in capturing data distributions over discrete sequences in domains such as text and proteins. These models generate data by iteratively unmasking tokens starting from a fully masked sequence, with the unmasking order typically chosen at random or using a heuristic based on denoiser probabilities. In this work, we propose a scheme for learning the unmasking order using an additional lightweight policy network on top of a diffusion model. Our proposed loss reweights terms in the masked diffusion loss according to policy probabilities, and results in a policy that prefers positions where the denoiser is more likely to be correct. We study this loss in two settings: (i) training solely the policy while using a frozen pre-trained denoiser, and (ii) training the policy and denoiser jointly with the weighted loss to allow for mutual adaptation. We demonstrate that our approach outperforms common heuristics on problems that are sensitive to token ordering, such as combinatorial tasks and proteins.

Related papers

• Stochastic MeanFlow Policies: One-Step Generative Control with Entropic Mirror Descent [53.828537014796574]
  Online off-policy reinforcement learning (RL) is shaped by two choices: the policy class and the update rule.<n>We propose MeanFlow Policies, a one-step generative policy class that maps noise to actions through a MeanFlow transformation.<n>Across seven MuJoCo benchmarks, S improves over Gaussian and generative baselines while retaining single-step inference efficiency.
  arXiv  Detail & Related papers  (2026-05-20T15:14:14Z)
• dUltra: Ultra-Fast Diffusion Language Models via Reinforcement Learning [36.12942468805232]
  Masked diffusion language models offer potential for parallel token generation.<n>Open-source MDLMs decode fewer than 5 tokens per model forward pass.<n>dUltra learns unmasking strategies for efficient parallel decoding.
  arXiv  Detail & Related papers  (2025-12-24T23:31:48Z)
• Learning Unmasking Policies for Diffusion Language Models [33.44995119635116]
  Language Models (dLLMs) now match the downstream performance of their autoregressive counterparts on many tasks.<n>One particularly successful variant is masked discrete diffusion, in which a buffer filled with special mask tokens is progressively replaced with tokens sampled from the model's vocabulary.<n>In this work, we propose to train sampling procedures using reinforcement learning.
  arXiv  Detail & Related papers  (2025-12-09T20:44:33Z)
• Improving Discrete Diffusion Unmasking Policies Beyond Explicit Reference Policies [47.6755955972232]
  We cast denoising as a KL-regularized Markov decision process (MDP) with an explicit reference policy and optimize a regularized objective.<n>We prove that the optimized policy under this framework generates samples that more closely match the data distribution than schedules.
  arXiv  Detail & Related papers  (2025-10-07T09:44:24Z)
• Continual-MAE: Adaptive Distribution Masked Autoencoders for Continual Test-Time Adaptation [49.827306773992376]
  Continual Test-Time Adaptation (CTTA) is proposed to migrate a source pre-trained model to continually changing target distributions. Our proposed method attains state-of-the-art performance in both classification and segmentation CTTA tasks.
  arXiv  Detail & Related papers  (2023-12-19T15:34:52Z)
• Learning to Mask and Permute Visual Tokens for Vision Transformer Pre-Training [55.12082817901671]
  We propose a new self-supervised pre-training approach, named Masked and Permuted Vision Transformer (MaPeT)<n>MaPeT employs autoregressive and permuted predictions to capture intra-patch dependencies.<n>Our results demonstrate that MaPeT achieves competitive performance on ImageNet, compared to baselines and competitors under the same model setting.
  arXiv  Detail & Related papers  (2023-06-12T18:12:19Z)
• Pre-training Co-evolutionary Protein Representation via A Pairwise Masked Language Model [93.9943278892735]
  Key problem in protein sequence representation learning is to capture the co-evolutionary information reflected by the inter-residue co-variation in the sequences. We propose a novel method to capture this information directly by pre-training via a dedicated language model, i.e., Pairwise Masked Language Model (PMLM) Our result shows that the proposed method can effectively capture the interresidue correlations and improves the performance of contact prediction by up to 9% compared to the baseline.
  arXiv  Detail & Related papers  (2021-10-29T04:01:32Z)
• Probabilistic fine-tuning of pruning masks and PAC-Bayes self-bounded learning [16.526326919313924]
  We study an approach to learning pruning masks by optimizing the expected loss of pruning masks. We analyze the training dynamics of the inducedadaptive predictor in the setting of linear regression. We show that a PAC-Bayes generalization error bound is controlled by the magnitude of the change in feature alignment between the 'prior' and 'posterior' data.
  arXiv  Detail & Related papers  (2021-10-22T14:25:22Z)
• Bi-Granularity Contrastive Learning for Post-Training in Few-Shot Scene [10.822477939237459]
  We propose contrastive masked language modeling (CMLM) for post-training to integrate both token-level and sequence-level contrastive learnings. CMLM surpasses several recent post-training methods in few-shot settings without the need for data augmentation.
  arXiv  Detail & Related papers  (2021-06-04T08:17:48Z)
• Improving Self-supervised Pre-training via a Fully-Explored Masked Language Model [57.77981008219654]
  Masked Language Model (MLM) framework has been widely adopted for self-supervised language pre-training. We propose a fully-explored masking strategy, where a text sequence is divided into a certain number of non-overlapping segments.
  arXiv  Detail & Related papers  (2020-10-12T21:28:14Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.