💬 NLP

A Continuous-Time Markov Chain Framework for Insertion Language Models

This paper establishes a continuous-time Markov chain framework to derive a principled diffusion-style denoising objective for Insertion Language Models, demonstrating that prior methods are special cases of this approach while achieving competitive performance and enhanced sampling flexibility.

Original authors: Dhruvesh Patel, Benjamin Rozonoyer, Soumitra Das, Tahira Naseem, Tim G. J. Rudner, Andrew McCallum

Published 2026-06-10

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Dhruvesh Patel, Benjamin Rozonoyer, Soumitra Das, Tahira Naseem, Tim G. J. Rudner, Andrew McCallum

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to write a story, but you have a very strict rule: you must write every word in order, from the very first letter to the very last. This is how most current AI language models work (called "Autoregressive" models). They are like a train that can only move forward on a single track. If the train gets stuck or makes a wrong turn early on, it's hard to fix the whole journey.

Other models try to be more flexible by guessing words in the middle of a sentence and filling in the blanks (like "Masked Diffusion Models"). However, these can sometimes struggle to figure out exactly how long the sentence should be or might get confused about the order of events.

This paper introduces a new way to teach AI to write, called Diffusion-based Insertion Language Models (DILMs). Here is how it works, using simple analogies:

The Core Idea: The "Scissors and Glue" Game

The authors imagine a game played in reverse to teach the AI how to write.

The Noising Process (The Scissors):
Imagine you have a perfect, complete sentence: "The quick brown fox jumps over the lazy dog."
The AI's teacher takes a pair of scissors and starts cutting out words one by one, at random.
- First, it cuts out "lazy."
- Then it cuts out "brown."
- Then "jumps."
- Eventually, you are left with just a few words or even an empty space.
  The paper uses a mathematical framework called a Continuous-Time Markov Chain to describe this cutting process. Think of this as a precise, scientific way of saying, "We will remove words at a steady, predictable rate until the sentence is gone."
The Learning Process (The Glue):
Now, the AI has to play the game in reverse. It starts with the empty space (or the few remaining words) and has to figure out how to put the words back in.
- Instead of just guessing the next word, the AI gets to insert words anywhere it wants.
- It can put "brown" between "The" and "quick."
- It can put "lazy" between "over" and "dog."
- It can even insert multiple words at once in different gaps.

Why is this special?

The paper claims this method combines the best of two worlds:

Flexibility: Unlike the "train on a track" model, this AI can fix mistakes or add details in the middle of a sentence without having to rewrite the whole thing from the start.
Knowing When to Stop: A common problem with these "fill-in-the-blank" models is that they don't know when the sentence is finished. They might keep adding words forever or stop too early.
- The authors solved this by teaching the AI to predict the "length-to-go." Imagine the AI has a little fuel gauge. As it inserts words, the fuel gauge goes down. When the gauge hits zero, the AI knows, "Okay, the sentence is complete," and it stops naturally.

The Two Versions of the New Model

The paper proposes two specific ways to do this "gluing" (insertion):

DILM-S (Single Insertion): The AI picks one spot and one word to insert at a time. It's like carefully placing one Lego brick at a time. This is very precise.
DILM-M (Multiple Insertions): The AI looks at all the empty spots at once and can fill in several gaps simultaneously. It's like grabbing a handful of Lego bricks and snapping them into place all at once. This is faster.

What did they find?

The researchers tested their new models on two types of tasks:

Planning Tasks (The "Star Graph" Game):
Imagine a map with a central hub and several paths leading out. The AI has to figure out the correct path from a start point to an end point.
- Result: The new models (DILM-S and DILM-M) were almost perfect at this. They were much better than the standard "train on a track" models and the other "fill-in-the-blank" models. They could see the whole picture and plan the route correctly.
Writing Stories (Language Modeling):
They tested the models on writing news articles and general text.
- Result: The new models were just as good as the best existing models at writing coherent text.
- The Big Bonus: The authors found that with their new method, you can trade off speed vs. quality. If you let the AI take more "steps" to insert words (more time), the writing gets better. If you rush it, it's faster but slightly less perfect. This gives users a knob to turn to get exactly what they need.

Summary

In short, this paper takes the messy, trial-and-error process of "filling in the blanks" and gives it a solid mathematical foundation. By treating the process like a continuous game of cutting and pasting words, they created an AI that can write in any order, knows exactly when to stop, and can be tuned to be either fast or very high-quality depending on the user's needs.

Technical Summary: A Continuous-Time Markov Chain Framework for Insertion Language Models

Problem Statement

Autoregressive language models (ARMs) generate sequences strictly from left-to-right, which limits their flexibility for tasks requiring non-sequential generation, such as planning and sub-goal based infilling. While Masked Diffusion Models (MDMs) allow for arbitrary generation orders and multi-token steps, they often struggle with variable-length sequences because they treat padding tokens as ordinary vocabulary items, impacting generation quality and restricting sampling to block-based left-to-right procedures.

Existing Insertion Language Models (ILMs) address these limitations by iteratively inserting tokens at arbitrary positions, naturally handling variable lengths and relative ordering constraints. However, current ILM formulations rely on ad-hoc training objectives and sampling procedures lacking a rigorous probabilistic foundation. Furthermore, existing ILMs often struggle with the challenge of determining when to stop generation.

Methodology

The authors propose Diffusion-based Insertion Language Models (DILMs), a principled framework derived from first principles using a Continuous-Time Markov Chain (CTMC) on the space of variable-length sequences.

1. The Noising Process

The forward noising process is formulated as a CTMC that deletes tokens from a ground-truth sequence $x_0$ uniformly at random over a time interval $[0, T]$ .

State Space: The state space includes both the sequence content and its length.
Deletion Mechanism: Tokens are dropped one at a time (or via independent rates) with a time-dependent rate $\sigma_t$ .
Two Variants:
- Joint Noising: Deletions occur as a single event where the number of tokens dropped follows a Poisson distribution, and positions are chosen uniformly.
- Independent Noising: Each position has an independent deletion rate, modeled via a Binomial distribution.

2. The Generative Process (Reverse CTMC)

The generative process is the time reversal of the noising process, constructing a sequence by inserting tokens. The authors derive an Evidence Lower Bound (ELBO) for the data log-likelihood under this reverse process. They propose two parameterizations for the reverse rate matrix $\hat{R}^\theta_t$ :

DILM-S (Single Insertion): Models the joint probability of an insertion location and a vocabulary item, as well as the "length-to-go" (the number of tokens remaining to be inserted). This parameterization predicts one insertion at a time.
DILM-M (Multiple Insertion): Models independent, per-position insertion rates. This allows the model to sample multiple insertions simultaneously across different gaps in a single step.

Both variants utilize a transformer backbone with adaptive layer normalization (AdaLN) conditioned on time $t$ . To handle the stopping condition, the models predict the length-to-go (or gap sizes), allowing the generation to terminate when the predicted insertion rate drops to zero, rather than relying on a separate stopping classifier or hyperparameters.

3. Sampling

Sampling is performed via Euler discretization (for DILM-S) or a constrained $\tau$ -leaping procedure (for DILM-M). These methods allow for a trade-off between sample quality and the number of forward passes (sampling steps), a flexibility not present in standard ILMs.

Key Contributions

Principled Framework: The paper derives a diffusion-style denoising objective for insertion-based generation from first principles using a CTMC framework, moving beyond ad-hoc formulations.
Unification: The proposed sampling procedures unify existing formulations of ILMs, showing them to be special cases of this broader framework.
Two Parameterizations: The authors introduce two variants (DILM-S and DILM-M) that differ in how they model insertion rates (joint vs. independent), offering flexibility in sampling efficiency and multi-token generation.
Theoretical Grounding: The work provides a rigorous derivation of the ELBO and discusses computational considerations for training and sampling.

Empirical Results

The authors evaluate DILMs on synthetic planning tasks and language modeling benchmarks.

Synthetic Planning (Star Graph Tasks): On variable-length planning tasks where ARMs and MDMs struggle, DILM-S and DILM-M achieve near-perfect accuracy (comparable to existing ILMs), demonstrating that the diffusion-based approach retains the benefits of insertion-based generation for complex structural tasks.
Language Modeling (LM1B and OpenWebText):
- On the LM1B benchmark, DILM-S and DILM-M are competitive with existing ILMs and MDMs in terms of per-token negative log-likelihood (NLL).
- On OpenWebText, DILM variants outperform both ILM and MDM baselines.
- Sampling Flexibility: A key finding is the ability to trade off generation quality against the number of sampling steps. Unlike existing ILMs, which show flat or erratic perplexity curves, DILMs exhibit a smooth improvement in generative perplexity as sampling steps increase.

Significance and Claims

The paper claims that DILMs offer a probabilistically principled alternative to existing insertion models, providing a rigorous training objective and sampling procedure. The framework retains the core advantages of ILMs—flexible generation order and natural variable-length handling—while addressing the limitations of ad-hoc stopping criteria and lack of theoretical grounding.

The authors note that while DILM-M allows for multiple insertions, it is currently constrained to one token per gap, limiting the modeling of contiguous token correlations. They identify multi-token span insertion and scaling to larger models as directions for future work. The work establishes a foundation for treating insertion-based generation as a continuous-time diffusion process, unifying previous heuristic approaches under a single theoretical umbrella.