Insertion Based Sequence Generation with Learnable Order Dynamics

Dhruvesh Patel; Benjamin Rozonoyer; Gaurav Pandey; Tahira Naseem; Ramón Fernandez Astudillo; Andrew McCallum

cs.LGFeb 21, 2026

Insertion Based Sequence Generation with Learnable Order Dynamics

Authors: Dhruvesh Patel, Benjamin Rozonoyer, Gaurav Pandey, Tahira Naseem, Ramón Fernandez Astudillo, Andrew McCallum

Organizations: University of Massachusetts Amherst · IBM Research

Abstract

Existing insertion-based masked diffusion models that generate sequences by interleaving token insertion with unmasking use fixed schedules that are not dependent on the data. For structured sequences like graphs and molecules, learning data-dependent generation orders can improve generation quality by reducing uncertainty over the action space. We propose LoFlexMDM, an insertion-based masked diffusion model with learnable order dynamics that learns data-dependent insertion and unmasking rates. We generalize the discrete flow matching framework to work with variable-length sequences, propose a tractable schedule parameterization and a training objective for joint training of the generator and the target order dynamics. On De Novo and fragment-constrained molecule generation, LoFlexMDM improves sample quality over FlexMDM by up to 17.5% and 6.7%, respectively. These results show that learning the target generation order can improve insertion-based diffusion models without giving up tractable training. We open source the code at https://github.com/dhruvdcoder/LoFlexMDM.

Figures & tables

Figure 1: Left: The left subfigure illustrates how the endpoint-dependent target process grows the final sequence

x_{1}

from the empty sequence (all [D] tokens) by sampling the insertion time

{T_{\color[rgb]{0.1484,0.2734,0.5664}\textnormal{in}}}^{i}\sim f^{\phi}_{{T_{\color[rgb]{0.1484,0.2734,0.5664}\textnormal{in}}}^{i}}(t;x_{1})

for the insertion transition

{\color[rgb]{0.1484,0.2734,0.5664}{{\textnormal{{[D]}}}}}\to{\color[rgb]{0.6016,0,0}{{\textnormal{{[M]}}}}}

, followed by the unmasking time

{T_{\color[rgb]{0.6016,0,0}\textnormal{um}}}^{i}\sim f^{\phi}_{{T_{\color[rgb]{0.6016,0,0}\textnormal{um}}}^{i}|{T_{\color[rgb]{0.1484,0.2734,0.5664}\textnormal{in}}}^{i}}(t;x_{1})

for the unmasking transition

{\color[rgb]{0.6016,0,0}{{\textnormal{{[M]}}}}}\to x_{1}^{i}

, for each position

i

. The target process induces a target generation order:

j

is generated after

i

{T_{\color[rgb]{0.6016,0,0}\textnormal{um}}}^{j}>{T_{\color[rgb]{0.6016,0,0}\textnormal{um}}}^{i}

. Middle: The middle subfigure (top), shows the auxiliary neural network

\phi

that takes in a clean sequence

x_{1}

and outputs per-token target insertion rates

\lambda_{{\color[rgb]{0.1484,0.2734,0.5664}{\textnormal{in}}}}^{\phi}

and unmasking rates

\lambda_{{\color[rgb]{0.6016,0,0}\textnormal{um}}}^{\phi}

, which induce the densities

f^{\phi}

. The map

f_{\text{contract}}

removes the [D] tokens from the intermediate

z_{t}

to obtain a partial sequence

x_{t}

. The generator network (bottom right) takes

x_{t}

and produces insertion

\lambda_{{\color[rgb]{0.1484,0.2734,0.5664}{\textnormal{in}}}}^{\theta}

and unmasking

\lambda_{{\color[rgb]{0.6016,0,0}\textnormal{um}}}^{\theta}\cdot K^{\theta}

rates which are trained to match the endpoint-dependent target rates

\lambda_{{\color[rgb]{0.1484,0.2734,0.5664}{\textnormal{in}}}}^{\phi}

and

\lambda_{{\color[rgb]{0.6016,0,0}\textnormal{um}}}^{\phi}

using Eq. (11) in Section 3.2. Right: The rightmost subfigure shows the generator network which consists of a transformer backbone, a linear LM head that produces token probabilities, and two rate heads that produce insertion and unmasking rates.

Figure 2: Projected Discrete Flow Matching

Figure 3: Kumaraswamy CDF shapes for different parameter values.

Panel of Figure 6 — Figure 6: Left : Quality vs sampling steps for de novo generation for Lo FlexMDM and FlexMDM. Right : Quality vs diversity for de novo generation for Lo FlexMDM and FlexMDM.

Model	$b_{\color[rgb]{0.6016,0,0}\textnormal{um}}$	$b_{\text{in}}$	conf.	$p$	Validity (%)	Diversity	Uniqueness (%)	Quality (%)
$\text{SAFE-GPT}^{\dagger}$					94.0 $\scriptstyle{\pm 0.4}$	0.879 $\scriptstyle{\pm 0.001}$	100.0 $\scriptstyle{\pm 0.0}$	54.7 $\scriptstyle{\pm 0.3}$
$\text{MDM}^{\ddagger}$					96.7 $\scriptstyle{\pm 0.3}$	0.896 $\scriptstyle{\pm 0.001}$	99.3 $\scriptstyle{\pm 0.2}$	53.8 $\scriptstyle{\pm 1.7}$
FlexMDM			×		98.9 $\scriptstyle{\pm 0.1}$	0.890 $\scriptstyle{\pm 0.000}$	99.6 $\scriptstyle{\pm 0.1}$	62.0 $\scriptstyle{\pm 0.7}$
FlexMDM			✓		67.8 $\scriptstyle{\pm 0.3}$	0.940 $\scriptstyle{\pm 0.000}$	61.7 $\scriptstyle{\pm 0.7}$	5.5 $\scriptstyle{\pm 0.4}$
Lo FlexMDM (Aux. Size=medium)			✓	1.0	99.0 $\scriptstyle{\pm 0.0}$	0.900 $\scriptstyle{\pm 0.000}$	99.3 $\scriptstyle{\pm 0.2}$	55.4 $\scriptstyle{\pm 0.3}$
			✓	0.5	99.9 $\scriptstyle{\pm 0.0}$	0.830 $\scriptstyle{\pm 0.000}$	93.2 $\scriptstyle{\pm 0.6}$	69.3 $\scriptstyle{\pm 0.7}$

Table 2: Results for de novo small molecule generation. For both FlexMDM and Lo FlexMDM, we use 1024 sampling steps. indicates frozen, indicates trainable; ✓/× indicates whether confidence-based position selection for unmasking is enabled, and

p

indicates the truncation threshold for nucleus sampling ( Holtzman et al., 2020 ) .

\dagger

Noutahi et al. (2024)

\ddagger

Lee et al. (2025)

Figure 7: Normalised generation time by SAFE token category for FlexMDM (left) and Lo FlexMDM (right). For each violin, the coloured horizontal line marks the mean and the black line marks the median.

Panel of Figure 8 — Figure 8: Example de novo generation trajectory. FlexMDM (left) and Lo FlexMDM (right) on the same molecule. Every line is a subsampled generation step; grey _ marks a still-masked position. Tokens are coloured by semantic category (legend above). The molecule graph below each trajectory colours atoms by normalised generation time ( blue = early, red = late). Lo FlexMDM first resolves ring closures and fragment separators, then places attachment brackets and aromatic atoms, and fills in aliphatic atoms, attachment labels, and bond/branch tokens last. FlexMDM shows no such systematic pattern. An additional example is in Appendix G.2.1 .

Model	Task	Val.	Div.	Uniq.	Qual.
$\text{ARM}^{\dagger}$	LD	76.6	0.545	82.5	21.7
	ME	96.1	0.562	66.8	18.6
	SD	97.7	0.575	74.7	10.0
	SG	95.7	0.573	83.0	14.3
$\text{MDM}^{\ddagger}$	LD	100.0	0.547	83.7	21.9
	ME	82.9	0.617	77.5	30.1

Table 3: Fragment-constrained molecule generation results. Means over 5 runs (Std. in Appendix 7 ). Tasks are abbreviated as LD (linker design), ME (motif extension), SD (scaffold decoration), and SG (superstructure generation). The best results are highlighted in bold and the second-best mean in each column is underlined. (

\dagger

Noutahi et al. (2024)

\ddagger

; Lee et al. (2025) )

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

General
Notation	Description
${\mathbb{X}}$ , ${\mathbb{Z}}$	Sets (blackboard font)
$[\![n]\!]$	Set of natural numbers $\{1,2,\ldots,n\}$
$x_{t}$	Particular state at time $t$ of a stochastic process $X_{t}$
$x_{t}^{i},s^{i},w^{i},$ etc.	$i$ -th component of $x,s,w$ respectively,i.e. superscript $i$ denotes the $i$ -th component in a sequence
$\delta_{x}(y)=\delta_{\{x\}}(y)$	Indicator function that is 1 if $x=y$ and 0 otherwise

Appendix

Table 4: Summary of notation used throughout the paper.

Figure 9: An example of a star graph with variable arm lengths from the hard subset of the star graph traversal task from Patel et al. (2025b) . The top section shows the string representation of the graph, where the input sequence consists of edges in a random order, followed by the starting and terminal nodes. The output part is the path (edges in the correct order) connecting the starting and terminal nodes. The figure is adapted from Patel et al. (2025b) with permission.

Figure 10: Comparison of truncated and rescaled Kumaraswamy CDFs for different parameter values.

Model	lr	wt decay	train_steps	sampling_steps	confidence
FlexMDM	$10^{-4}$	$10^{-2}$	80,000	500	×
Lo FlexMDM	$10^{-4}$	$10^{-2}$	80,000	500	×

Appendix

Table 5: Hyperparameters for star graphs experiments.

Model	lr	wt decay	train_steps	sampling_steps	confidence
FlexMDM	$10^{-4}$	$10^{-2}$	50,000	1024	×
LoFlexMDM	$10^{-4}$	$10^{-2}$	50,000	1024	✓

Appendix

Table 6: Hyperparameters for star graphs experiments.

Model	Task	Validity (%)	Diversity	Uniqueness (%)	Quality (%)
$\text{ARM}^{\dagger}$	Linker design	76.6 $\scriptstyle{\pm 5.1}$	0.545 $\scriptstyle{\pm 0.007}$	82.5 $\scriptstyle{\pm 1.9}$	21.7 $\scriptstyle{\pm 1.1}$
	Motif extension	96.1 $\scriptstyle{\pm 1.9}$	0.562 $\scriptstyle{\pm 0.003}$	66.8 $\scriptstyle{\pm 1.2}$	18.6 $\scriptstyle{\pm 2.1}$
	Scaffold decoration	97.7 $\scriptstyle{\pm 0.3}$	0.575 $\scriptstyle{\pm 0.008}$	74.7 $\scriptstyle{\pm 2.5}$	10.0 $\scriptstyle{\pm 1.4}$
	Superstructure generation	95.7 $\scriptstyle{\pm 2.0}$	0.573 $\scriptstyle{\pm 0.028}$	83.0 $\scriptstyle{\pm 5.9}$	14.3 $\scriptstyle{\pm 3.7}$
$\text{MDM}^{\ddagger}$	Linker design	100.0 $\scriptstyle{\pm 0.0}$	0.547 $\scriptstyle{\pm 0.002}$	83.7 $\scriptstyle{\pm 0.5}$	21.9 $\scriptstyle{\pm 0.4}$
	Motif extension	82.9 $\scriptstyle{\pm 0.1}$	0.617 $\scriptstyle{\pm 0.002}$	77.5 $\scriptstyle{\pm 0.1}$	30.1 $\scriptstyle{\pm 0.4}$

Appendix

Table 7: Fragment-constrained molecule generation results (full). Means and standard deviations over 3 runs. The best results are highlighted in bold; the second-best mean in each column is underlined. (

\dagger

Noutahi et al. (2024)

\ddagger

; Lee et al. (2025) )

Figure 11: Loss comparison for the hard star graph traversal task with and without fixing target unmasking rate (

{\color[rgb]{0.6016,0,0}b_{{\color[rgb]{0.6016,0,0}\textnormal{um}}}}=1

). The erratic dynamics of a fully unconstrained set of target rates motivates freezing

b_{\color[rgb]{0.6016,0,0}\textnormal{um}}

Figure 12: Correlation between normalized generation order and distance from the junction node ( Fig. 5(b) ), shown at two-column width for clarity. Only trajectories with 100% exact match are included.

Figure 13: Correlation between normalized generation order and distance from the junction node, without filtering to 100% exact-match trajectories.

Figure 14: Filtered (100% exact match) vs. unfiltered generation-order correlation for each sampling setting.

Task	Model	Exact Match (%)	Token Accuracy (%)
star_medium	ARM	75.0	81.4
star_medium	MDM	36.5	90.6
	FlexMDM
	(conf=null)	89.6	95.7
	(conf=top_prob)	91.3	97.0
	Lo FlexMDM (Aux. Size=small)

Appendix

Table 8: Raw results for star graph traversal tasks.

Category	FlexMDM	Lo FlexMDM	$\Delta\mu$
Ring closure	$.71\pm.26$	$.48\pm.28$	$+.23$
Frag sep ( . )	$.67\pm.27$	$.55\pm.18$	$+.12$
Aromatic atom	$.65\pm.26$	$.55\pm.24$	$+.10$
Attach bracket	$.63\pm.27$	$.61\pm.19$	$+.02$
Attach label	$.68\pm.27$	$.80\pm.24$	$-.12$
Aliphatic atom	$.68\pm.27$	$.80\pm.16$	$-.13$

Appendix

Table 9: Summary of mean normalised generation time per token category (

\pm

std).

\Delta\mu

is the difference (FlexMDM

-

Lo FlexMDM), where positive values indicate FlexMDM generates that category later.

Panel of Figure 16 — Figure 16: Additional de novo generation trajectory. FlexMDM (left) and Lo FlexMDM (right) on the same molecule. Every line is a subsampled generation step; grey _ marks a still-masked position. Tokens are coloured by semantic category (legend above). The molecule graph below each trajectory colours atoms by normalised generation time ( blue = early, red = late). Lo FlexMDM first resolves ring closures and fragment separators, then places attachment brackets and aromatic atoms, and fills in aliphatic atoms, attachment labels, and bond/branch tokens last. FlexMDM shows no such systematic pattern.

Model					Validity (%)	Diversity	Uniqueness (%)	Quality (%)
	$b_{\color[rgb]{0.6016,0,0}\textnormal{um}}$	$b_{\text{in}}$	conf.	$p$
$\text{SAFE-GPT}^{\dagger}$					94.0 $\scriptstyle{\pm 0.4}$	0.879 $\scriptstyle{\pm 0.001}$	100.0 $\scriptstyle{\pm 0.0}$	54.7 $\scriptstyle{\pm 0.3}$
$\text{MDM}^{\ddagger}$					96.7 $\scriptstyle{\pm 0.3}$	0.896 $\scriptstyle{\pm 0.001}$	99.3 $\scriptstyle{\pm 0.2}$	53.8 $\scriptstyle{\pm 1.7}$
FlexMDM
			✓	1.0	67.8 $\scriptstyle{\pm 0.3}$	0.940 $\scriptstyle{\pm 0.000}$	61.7 $\scriptstyle{\pm 0.7}$	5.5 $\scriptstyle{\pm 0.4}$
			✓	0.9	67.3 $\scriptstyle{\pm 0.9}$	0.930 $\scriptstyle{\pm 0.000}$	56.9 $\scriptstyle{\pm 0.8}$	5.3 $\scriptstyle{\pm 0.2}$

Appendix

Table 10: Full set of results for de novo small molecule generation. For both FlexMDM and Lo FlexMDM, we use 1024 sampling steps. indicates frozen, indicates trainable; ✓/× indicates whether confidence-based position selection for unmasking is enabled, and

p

indicates the truncation threshold for nucleus sampling ( Holtzman et al., 2020 ) .

\dagger

Noutahi et al. (2024)

\ddagger

Lee et al. (2025)

Diffusion Models

Diffusion models learn to reverse gradual data corruption, enabling flexible generation, editing, forecasting, and inference.

See all 1,169 papers

Birth, Death and Diff: Insertion-Based Sequence Generation with Corrections
Oct 25, 2026Durga Prasad Maram, Dhruvesh Patel, Tahira Naseem +4Diffusion Language Models Text Generation
University of Massachusetts Amherst · IBM Research
Variational Learning for Insertion-based Generation
Jun 1, 2026Yangtian Zhang, Zhe Wang, Arthur Gretton +4Non-Autoregressive Text Generation Goal-Conditioned Planning
1Yale University · 2Google DeepMind · University College London +1
Learning Generation Orders for Masked Discrete Diffusion Models via Variational Inference
Feb 27, 2026David Fox, Sam Bowyer, Song Liu +3Parallel Decoding for Diffusion Language Models Masked Diffusion Models
Mistral AI

Oct 25, 2026J/K move · Enter open · S save

Birth, Death and Diff: Insertion-Based Sequence Generation with Corrections

Durga Prasad Maram, Dhruvesh Patel, Tahira Naseem, Benjamin Rozonoyer+3

University of Massachusetts Amherst · IBM Research

Diffusion Language Models Text Generation Diffusion Language Model Inference Parallel Decoding for Diffusion Language Models

Open paper AskPDF arXiv

Insertion-based masked diffusion models enable flexible-length, parallel generation by interleaving mask insertion with token unmasking. However, because their generation decisions are irreversible, aggressive parallel decoding can introduce errors that subsequent steps cannot correct, necessitating conservative decoding with more forward passes (NFEs). Moreover, they exhibit higher entropy than standard MDMs because they marginalize over multiple alignments during training and therefore require even more conservative decoding in the absence of a correction mechanism. To address this limitation, we introduce

δ

MDM, which makes intermediate generations revisable by equipping an insertion-based model with substitution and deletion operations. During training, we expose the model to errors caused by aggressive unmasking and insertion, teaching it to correct its own mistakes. At inference time, we interleave insertion and unmasking with correction steps that can revise previously unmasked tokens and remove unnecessary masks. This mechanism makes aggressive decoding practical, improving both generation speed and quality.

δ

MDM Pareto-dominates FlexMDM, a vanilla insertion-based MDM, on the speed--quality trade-off. On high-diversity unconditional text generation on OpenWebText, it halves the generative perplexity relative to FlexMDM at the highest inference budget. It also achieves lower generative perplexity at all inference budgets compared to MDM trained from scratch with variable-length padded sequences. On low-diversity conditional text generation on TinyGSM, it matches FlexMDM accuracy with half as many NFEs.

Abstract

Diffusion Models

Explore similar work

Birth, Death and Diff: Insertion-Based Sequence Generation with Corrections

Variational Learning for Insertion-based Generation

Learning Generation Orders for Masked Discrete Diffusion Models via Variational Inference