Any-Order Flexible Length Masked Diffusion

2025-08-31

AI Summary

Masked diffusion models (MDMs) are constrained by fixed-length generation, limiting their applicability to flexible sequence modeling. To address this, we propose FlexMDMs—the first diffusion model supporting variable-length discrete sequence generation. FlexMDMs enable non-autoregressive, parallel inference with arbitrary sequence lengths by dynamically inserting mask tokens and permitting denoising in arbitrary orders. Theoretically, we generalize the stochastic interpolation framework to accommodate both length variation and arbitrary denoising schedules. Engineering-wise, FlexMDMs seamlessly integrate with pre-trained MDMs, requiring only lightweight fine-tuning for adaptation. Empirically, FlexMDMs achieve ~60% higher success rates on maze planning; after three-day fine-tuning of LLaDA-8B, they outperform baselines significantly on GSM8K (58% → 67%) and code completion (52% → 65%), demonstrating both efficacy and practicality.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: Generation

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Masked diffusion models (MDMs) have recently emerged as a promising alternative to autoregressive models over discrete domains. MDMs generate sequences in an any-order, parallel fashion, enabling fast inference and strong performance on non-causal tasks. However, a crucial limitation is that they do not support token insertions and are thus limited to fixed-length generations. To this end, we introduce Flexible Masked Diffusion Models (FlexMDMs), a discrete diffusion paradigm that simultaneously can model sequences of flexible length while provably retaining MDMs' flexibility of any-order inference. Grounded in an extension of the stochastic interpolant framework, FlexMDMs generate sequences by inserting mask tokens and unmasking them. Empirically, we show that FlexMDMs match MDMs in perplexity while modeling length statistics with much higher fidelity. On a synthetic maze planning task, they achieve $approx 60 %$ higher success rate than MDM baselines. Finally, we show pretrained MDMs can easily be retrofitted into FlexMDMs: on 16 H100s, it takes only three days to fine-tune LLaDA-8B into a FlexMDM, achieving superior performance on math (GSM8K, $58% o 67%$) and code infilling performance ($52% o 65%$).

Problem

• Overcoming fixed-length generation in masked diffusion models
• Enabling token insertions for flexible sequence lengths
• Retaining any-order inference flexibility while modeling length variability

Innovation

• Flexible Masked Diffusion Models for variable length
• Inserting mask tokens and unmasking them
• Retrofitting pretrained MDMs into FlexMDMs via fine-tuning