Masked Language Flow Models

Abstract: Masked Diffusion Models (MDMs) promise fast, parallel language generation, but their reverse transition factorises across token positions -- an approximation that breaks down in the few-step sampling regime where parallel generation ought to provide the greatest efficiency gains. Flow LLMs (FLMs) sidestep this limitation by learning a continuous flow that transports noise toward clean sequences represented in Euclidean space, inducing a flow map that can be distilled for single-step generation. However, this makes complex tasks requiring multi-step reasoning problematic for FLMs, as FLMs are forced to decode every token during generation. To address this, we introduce Masked Language Flow Models (MLFMs), which incorporate masking into FLMs using a continuous stochastic interpolant to bridge partially masked and clean sequences. This design enables conditional generation via continuous flows and allows pretrained MDMs to be converted into MLFMs through a simple, lightweight adaptation. Leveraging this flexibility, we propose a novel sampler that alternates continuous denoising with the discrete unmasking of confident tokens to better support multi-step reasoning. We evaluate our approach on GSM8K and MT-Bench and find, for the first time, that flow-based LLMs can be scaled to solve downstream reasoning and instruction-following tasks.

• The paper introduces Masked Language Flow Models (MLFMs), a hybrid that unifies masked diffusion and flow-based modeling for conditional generation.
• It employs a Brownian bridge and context-corrupted classifier-free guidance to seamlessly transition from masked to clean token embeddings.
• Empirical results show improved instruction-following performance and reduced sampling steps, paving the way for scalable, parallel generation.

Masked Language Flow Models: Integrated Conditional Flow-Based Generation

Conceptual Framework and Model Architecture

Masked Language Flow Models (MLFMs) are introduced as a hybrid architecture combining the strengths of Masked Diffusion Models (MDMs) and Flow LLMs (FLMs). MLFMs utilize a Brownian bridge stochastic interpolant to seamlessly connect partially masked sequences with fully clean sequences in continuous Euclidean space. This design enables the model to retain any-position conditional generation from MDMs while leveraging the smooth joint token evolution of FLMs.

The MLFM formalism defines a forward process $z_{s,t}$, parameterized by masking probability $s$ and noising time $t$. Clean token positions remain fixed in their embeddings, whereas masked positions evolve continuously according to a Brownian bridge from the mask embedding $m$ to the token embedding $E_{X^\ell}$. At $t=0$, masked positions are deterministically set to $m$; at $t=1$, they arrive at $E_{X^\ell}$, ensuring endpoint equivalence with the discrete mask representation.

MLFMs employ a denoiser $p_\theta$ to predict token distributions at masked positions, trained via cross-entropy only on masked tokens. This loss function is a continuous extension of the typical MDM objective and enables adaptation of pretrained MDMs to the MLFM architecture with minimal compute overhead by reusing the token embeddings and transformer backbone, supplementing them with LoRA adapters and AdaLN modules for time conditioning.

Sampling and Conditional Generation

MLFMs facilitate a novel sampling strategy that combines flow-based parallel generation with online token unmasking, exploiting classifier-free guidance. High-confidence predictions (mode probability $s$0) are promoted immediately to clean embeddings, allowing subsequent denoising steps to condition on resolved tokens rather than corrupted states. This contrasts with traditional diffusion sampling, where latent variable positions remain noisy until terminal time.

Classifier-free guidance is adapted via context-corrupted guidance: two parallel model calls are performed, one conditioning on promoted clean tokens, the other corrupting these embeddings according to the forward bridge marginal. The difference isolates model sensitivity to clean context, enhancing guidance without violating forward process support.

Theoretical analysis bounds the total variation distance between the promoted sampler and the true distribution by $s$1, guaranteeing negligible error for small $s$2 and providing signal amplification analogous to self-conditioning, previously unavailable in hybrid discrete-continuous LLMs.

Figure 1: MT-Bench and GSM8K performance as functions of guidance scale and sampler steps, highlighting optimal guidance parameters and step tradeoffs.

Empirical Results and Ablations

MLFM adaptation is conducted on a 1028M parameter SMDM backbone, with the architectural modifications described above. Training is performed in two phases: continuous adaptation on SlimPajama corpus and instruction-centric supervised fine-tuning on ShareGPT, MetaMathQA, GSM8K-Aug-NL, NuminaMath-CoT, and OpenCodeInstruct datasets.

On downstream reasoning and instruction-following benchmarks, MLFM demonstrates notable strengths. On MT-Bench (instruction-following via GPT-4o scoring), MLFM achieves a first-turn score of 2.27, significantly exceeding SMDM (1.60) and autoregressive baseline (1.57), using only 128 sampling steps versus SMDM's 256 steps. On GSM8K (math reasoning), MLFM attains 31.24% accuracy, trailing AR and SMDM baselines; this is attributed to differences in fine-tuning protocol rather than modeling fundamental limitations.

Ablations confirm that both context-correlated classifier-free guidance and online token promotion substantially improve performance, particularly when using larger guidance scales and increased sampling steps. Performance gains saturate with step count, and strong guidance is especially impactful for instruction tasks.

Practical and Theoretical Implications

MLFMs advance flow-based language modeling by supporting exact conditional generation in continuous space with efficient training/finetuning pathways. The architecture enables scalable, parallel generation while retaining multi-step reasoning (tokens can be promoted as soon as resolved), addressing restrictive assumptions in prior FLMs and MDMs.

Practically, MLFM offers the promise of more efficient inference than strictly autoregressive models, making it viable for long-sequence or high-throughput applications. The adaptability of pretrained MDMs suggests a path for leveraging large existing models and distilling flow maps for further acceleration.

Theoretically, MLFMs open avenues for hybrid conditional generation architectures that couple discrete and continuous dynamics, facilitating research into adaptive promotion rules, bridge distributions, and guidance mechanisms. Future developments may include distillation strategies, scaling to larger parameter counts, and extending to code generation or other modalities.

Conclusion

Masked Language Flow Models unify mask-based and flow-based modeling paradigms via a Brownian bridge in embedding space, enabling flexible, conditional, parallel generation with robust sampling schemes. The architecture supports efficient adaptation from existing mask diffusion models and demonstrates strong performance on conditional tasks requiring reasoning and instruction-following. MLFMs provide a foundation for further research into scalable flow-based generation, distillation, and hybrid modeling protocols.