Learning Generation Orders for Masked Discrete Diffusion Models via Variational Inference

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Fox, David, Bowyer, Sam, Liu, Song, Aitchison, Laurence, Santos-Rodriguez, Raul, Yang, Mengyue
Natura: Preprint
Pubblicazione: 2026
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866915823250046976
author Fox, David
Bowyer, Sam
Liu, Song
Aitchison, Laurence
Santos-Rodriguez, Raul
Yang, Mengyue
author_facet Fox, David
Bowyer, Sam
Liu, Song
Aitchison, Laurence
Santos-Rodriguez, Raul
Yang, Mengyue
contents Masked discrete diffusion models (MDMs) are a promising new approach to generative modelling, offering the ability for parallel token generation and therefore greater efficiency than autoregressive counterparts. However, achieving an optimal balance between parallel generation and sample quality remains an open problem. Current approaches primarily address this issue through fixed, heuristic parallel sampling methods. There exist some recent learning based approaches to this problem, but its formulation from the perspective of variational inference remains underexplored. In this work, we propose a variational inference framework for learning parallel generation orders for MDMs. As part of our method, we propose a parameterisation for the approximate posterior of generation orders which facilitates parallelism and efficient sampling during training. Using this method, we conduct preliminary experiments on the GSM8K dataset, where our method performs competitively against heuristic sampling strategies in the regime of highly parallel generation. For example, our method achieves 33.1\% accuracy with an average of only only 4 generation steps, compared to 23.7-29.0\% accuracy achieved by standard competitor methods in the same number of steps. We believe further experiments and analysis of the method will yield valuable insights into the problem of parallel generation with MDMs.
format Preprint
id arxiv_https___arxiv_org_abs_2602_23968
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Learning Generation Orders for Masked Discrete Diffusion Models via Variational Inference
Fox, David
Bowyer, Sam
Liu, Song
Aitchison, Laurence
Santos-Rodriguez, Raul
Yang, Mengyue
Machine Learning
Masked discrete diffusion models (MDMs) are a promising new approach to generative modelling, offering the ability for parallel token generation and therefore greater efficiency than autoregressive counterparts. However, achieving an optimal balance between parallel generation and sample quality remains an open problem. Current approaches primarily address this issue through fixed, heuristic parallel sampling methods. There exist some recent learning based approaches to this problem, but its formulation from the perspective of variational inference remains underexplored. In this work, we propose a variational inference framework for learning parallel generation orders for MDMs. As part of our method, we propose a parameterisation for the approximate posterior of generation orders which facilitates parallelism and efficient sampling during training. Using this method, we conduct preliminary experiments on the GSM8K dataset, where our method performs competitively against heuristic sampling strategies in the regime of highly parallel generation. For example, our method achieves 33.1\% accuracy with an average of only only 4 generation steps, compared to 23.7-29.0\% accuracy achieved by standard competitor methods in the same number of steps. We believe further experiments and analysis of the method will yield valuable insights into the problem of parallel generation with MDMs.
title Learning Generation Orders for Masked Discrete Diffusion Models via Variational Inference
topic Machine Learning
url https://arxiv.org/abs/2602.23968