Insertion Based Sequence Generation with Learnable Order Dynamics

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.

Generative models often create sequences one piece at a time from left-to-right, but structured sequences such as molecules and graph paths do not always have one natural left-to-right order. For these sequences, the order in which parts are created can make the task easier or harder. Insertion-based models can build variable-length sequences by adding pieces in different places, but they usually follow random generation orders that ignore the structure of each example. We introduce LFlexMDM, a model that learns when and where to add new pieces while it generates a sequence. Instead of forcing the model to consider many possible orders equally, LFlexMDM learns generation orders that fit the data. For graph paths, it learns to grow paths from endpoints toward important junctions. For molecules, it learns to resolve structural features such as rings, fragments, and attachment points before filling in finer details. This learned ordering improves the quality of generated molecular sequences in both standard and fragment-constrained molecule generation tasks. The result suggests that teaching generative models how to choose a good construction order can make them better at creating structured scientific sequences.