🤖 machine learning

Variational Learning for Insertion-based Generation

This paper introduces the Insertion Process (IP), a variational learning framework that enables variable-length non-monotonic sequence generation by learning data-driven insertion orders through a bijective correspondence between insertion trajectories and permutations, thereby improving modeling quality and generalization in domains lacking a canonical left-to-right structure.

Original authors: Yangtian Zhang, Zhe Wang, Arthur Gretton, Rex Ying, David van Dijk, Michalis K. Titsias, Jiaxin Shi

Published 2026-06-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yangtian Zhang, Zhe Wang, Arthur Gretton, Rex Ying, David van Dijk, Michalis K. Titsias, Jiaxin Shi

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 build a complex Lego castle.

The Old Way (Autoregressive Models):
Most AI models today build like a strict construction crew that only works from left to right. They lay the first brick, then the second, then the third, and so on. If they realize halfway through that they need a tower in the middle of the wall, they can't just pop a brick in there. They have to tear everything down and start over, or awkwardly try to build the tower on top of the existing wall, which often leads to a wobbly, unstable mess. This works fine for simple sentences, but it struggles with complex structures where the "middle" depends on the "end."

The Problem with Current "Flexible" Models:
Scientists have tried to fix this with "masked diffusion" models. Think of these as a crew that starts with a pre-determined number of empty Lego slots (a fixed grid) and randomly fills them in.

The Flaw 1: They are stuck with a fixed number of slots. If you want a tiny house or a massive skyscraper, you have to guess the size beforehand.
The Flaw 2: They don't really care when they fill a slot. They might fill the roof before the foundation, or the chimney before the walls. While this is flexible, it's inefficient because the AI has to learn to be consistent with every possible order of construction, which is a massive, confusing task.

The New Solution: The "Insertion Process" (IP)
The authors of this paper introduce a new way to build called the Insertion Process. Imagine a master builder who doesn't follow a fixed grid or a left-to-right rule. Instead, they have a magical ability to:

Look at the half-built structure.
Decide exactly where to insert the next piece (e.g., "I'll put this window right in the middle of the wall").
Decide what piece to put there.
Decide when to stop (e.g., "The castle is done, no more bricks needed").

This model learns to build in the best order for the specific task, rather than forcing a rigid order.

How They Made It Work (The "Magic Trick")
The tricky part was teaching the AI to learn this "best order" mathematically. If you try to teach it by letting it guess random insertion spots, the math gets messy and unstable.

The authors discovered a clever mathematical shortcut. They realized that every possible way to build a sequence is actually just a different way of ordering the same final pieces.

Analogy: Imagine you have a deck of cards. Whether you deal them one by one from the top, or shuffle them and deal them in a weird pattern, the final hand is the same. The "Insertion Process" translates the messy act of "inserting into a growing pile" into a clean, organized "permutation" (a specific ordering) of the final cards.
By doing this translation, they could use a standard, efficient training method (Variational Inference) to teach the AI which ordering is best for the data it's seeing.

What They Tested It On
The paper tested this new builder on two specific challenges where the "left-to-right" rule makes no sense:

Planning a Path (Mazes):
- The Task: The AI had to generate a path through a maze, visiting specific checkpoints.
- The Result: The old "left-to-right" builders got stuck because they couldn't easily jump back to fill a gap in the path. The new Insertion Process could jump around, filling in the gaps between checkpoints perfectly. It solved almost 100% of the hardest mazes, while the old methods failed miserably.
Designing Molecules (Chemistry):
- The Task: The AI had to generate valid chemical strings (SMILES). These strings are like a map of a molecule. A valid map needs parentheses to match up (like opening and closing a ring) and numbers to connect distant parts.
- The Result: The new model learned a "smart" way to build. It didn't just add atoms randomly. It learned to first build the "skeleton" (the rings and connections) and then fill in the atoms.
- Why it matters: This allowed the AI to create valid, unique, and novel molecules much better than models that just append pieces to the end. It learned the "grammar" of chemistry without being explicitly told the rules.

The Bottom Line
This paper presents a new way for AI to generate sequences (like text, plans, or chemical formulas) that doesn't force a rigid "start-to-finish" order. Instead, it learns to insert pieces exactly where they are needed, when they are needed. By using a clever math trick to translate "insertion" into "ordering," they trained a model that is better at handling complex, variable-length tasks than previous methods.

Technical Summary: Variational Learning for Insertion-based Generation

1. Problem Statement

Autoregressive sequence models typically generate tokens in a fixed left-to-right order. While effective for many tasks, this factorization is not inherent to all data structures and can be misaligned with real-world generation problems such as biological sequence design, program synthesis, and structured planning, where output parts may depend on future context.

Existing alternatives, such as masked diffusion models (non-monotonic generation), offer flexibility by updating subsets of masked positions on a fixed-length grid. However, these approaches suffer from two primary limitations:

Static Canvas: They operate on prescribed, fixed-length grids, requiring external length specifications and making variable-length generation indirect.
Order Agnosticism: Their training objectives often marginalize over or implicitly assume many possible update schedules. This forces the model to approximate conditionals for a factorial number of generation orders, even if only a small subset aligns with the data's underlying structure.

Recent work has addressed these issues in isolation: insertion-based models allow variable-length sequences via arbitrary insertions, while learning-order methods introduce distributions over permutations. However, combining these is non-trivial because in insertion models, insertion locations are not fixed coordinates but evolve with the partial sequence, tightly coupling generation order and intermediate states.

2. Methodology: The Insertion Process (IP)

The authors propose the Insertion Process (IP), a probabilistic framework that jointly learns:

Where to insert: A data-driven preference over insertion decisions (slot selection).
What to insert: A token distribution.
When to terminate: A stopping rule (policy-based or classifier-based).

2.1. Reparametrization via Permutations

A core theoretical contribution is the formalization of a bijective correspondence between insertion trajectories and permutations of the target sequence indices.

The Challenge: Direct variational inference over insertion trajectories is difficult because relative slot indices ( $z_i$ ) are trajectory-dependent, making it hard to define a stable posterior over order in the space of relative actions.
The Solution: The authors map every valid insertion trajectory producing a sequence $y_L$ to a unique permutation $\sigma \in S_L$ , where $\sigma_i$ is the index of the element filled at step $i$ .
Deterministic Mapping: The relative insertion slot at step $i$ becomes a deterministic function of the permutation prefix ( $\sigma_{\le i}$ ). Specifically, $z_i = 1 + |\{j \in \sigma_{<i} : j < \sigma_i\}|$ .
Likelihood Factorization: This change of variables converts the marginalization over trajectory-dependent actions into an exact sum over permutations, enabling an exact reparameterization of the data likelihood:
$p(y_L) = \sum_{\sigma \in S_L} p(y_L, \sigma)$

2.2. Variational Inference and Training

To train the model, the authors maximize an Evidence Lower Bound (ELBO) on the data log-likelihood, treating the generation order (permutation $\sigma$ ) as a latent variable.

Variational Posterior: They introduce an amortized variational posterior $q_\theta(\sigma | y_L)$ , parameterized as a Plackett-Luce (PL) model.
Optimization: Since the latent order is discrete, the objective is non-differentiable. The authors optimize the ELBO using:
- Gumbel-Top-k Trick: To sample permutations efficiently from the PL distribution.
- REINFORCE with Leave-One-Out (RLOO): To estimate gradients.
- Variance Reduction: By exploiting the permutation-trajectory structure, the inner expectation over the next index $\sigma_i$ (conditioned on a fixed prefix) is evaluated exactly in closed form, reducing the variance of the Monte Carlo estimator.

2.3. Network Architecture

Generative Decoder: A Transformer-based decoder processes the partial sequence (augmented with boundary tokens) to produce contextual embeddings.
- A location head scores candidate insertion slots (including a termination action) based on embeddings.
- A content head predicts the token distribution conditioned on the selected slot.
Inference Encoder: A bidirectional Transformer encoder processes the full observed sequence to output scores for each global index, which parameterize the Plackett-Luce distribution for sampling the generation order.

3. Key Contributions

Probabilistic Framework for Variable-Length Insertion: The first framework to learn insertion order for variable-length sequences using a bijective mapping to permutations, enabling exact likelihood reparameterization.
Permutation-Marginalized Likelihood: A theoretical derivation showing that the data likelihood can be expressed as a sum over permutations, facilitating variational inference with a latent order variable.
Efficient Training Algorithm: A variance-reduced training method using Plackett-Luce models, Gumbel-Top-k sampling, and RLOO gradients to handle the discrete, non-differentiable nature of the latent order.
Native Variable-Length Support: Unlike fixed-canvas diffusion models, IP natively supports variable-length generation and learns data-driven preferences for insertion orders.

4. Experimental Results

The authors evaluated IP on two domains: goal-conditioned planning and molecular string generation.

4.1. Goal-Conditioned Planning

Benchmarks: Maze planning (Braided, Imperfect, Perfect) and Star Graph planning.
Baselines: Fixed-Order ARMs (FO-ARM), Any-Order ARMs (AO-ARM), Masked Diffusion Models (MDM), and Learning-Order ARMs (LO-ARM).
Results: IP achieved state-of-the-art performance across all maze families and graph benchmarks.
- On Perfect Mazes (highly constrained), FO-ARMs degraded significantly, while IP maintained near-perfect accuracy (e.g., 99.9% on Hard Perfect Mazes vs. 27.2% for FO-ARM).
- IP substantially improved sequence and token accuracy and reduced edit distances compared to all baselines, demonstrating that learning instance-specific insertion orders is crucial for tasks with long-range constraints.

4.2. Molecular String (SMILES) Generation

Benchmark: GuacaMol unconditional and conditional generation.
Results:
- Validity & Novelty: IP achieved high validity (97.2–97.4%) and the highest novelty (~95%) among all models, significantly outperforming random-order insertion (AO-IP) which dropped to 88.2% validity. This indicates IP successfully internalizes chemical syntax rules.
- Distributional Metrics: While FO-ARMs slightly outperformed IP on distributional similarity (KL divergence, FCD), IP offered a superior trade-off by exploring the chemical space more effectively.
- Conditional Generation: In tasks requiring insertion between fixed fragments (e.g., linker design), IP maintained ~100% validity, whereas fixed-order models struggle with such constraints.
Interpretability: Analysis of the learned insertion schedules revealed an interpretable, consistent generation order: (1) building the molecular skeleton (rings/parentheses), (2) inserting atoms, and (3) terminating. This occurred without explicit grammatical priors.

5. Significance and Claims

The paper claims that learning data-dependent insertion orders significantly improves modeling quality and generalization in domains lacking a canonical left-to-right structure. By formalizing the connection between insertion trajectories and permutations, the authors provide a tractable, unbiased variational method for training insertion-based models.

The authors position their work as a step toward generative modeling in domains where the "natural" generation order is ambiguous or non-linear. They explicitly state that future work could apply this method to other domains such as protein design. The paper maintains a modest tone, noting that while IP improves generalization and validity in complex structural tasks, it does not necessarily outperform fixed-order models on simple distributional matching metrics, highlighting a trade-off between flexibility and strict distributional fidelity.