🤖 machine learning

Generative Molecular Morphing for Flexible-Size Design via Unbalanced Optimal Transport

This paper introduces Morph, a flexible-size generative model based on unbalanced optimal transport that overcomes the fixed-atom limitations of current diffusion methods to enable superior property steering and out-of-distribution 3D molecular design.

Original authors: Malte Franke, Stefan P. Schmid, Zarko Ivkovic, Kjell Jorner, Andreas Krause

Published 2026-06-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Malte Franke, Stefan P. Schmid, Zarko Ivkovic, Kjell Jorner, Andreas Krause

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 a master architect trying to design new buildings (molecules) for a city. In the past, the most advanced architects (AI models) had a strict rule: they could only design buildings with a fixed number of rooms.

If they wanted to build a house, they had to decide beforehand: "This house will have exactly 4 rooms." If they wanted a skyscraper, they had to decide: "This skyscraper will have exactly 50 rooms." They couldn't start with a sketch of a 2-room cabin and magically expand it into a mansion, nor could they shrink a mansion down to a cottage. They were stuck with the room count they started with.

This paper introduces a new AI architect named Morph. Morph is special because it doesn't care about the number of rooms. It can start with a blank slate (or a small sketch) and dynamically add or remove rooms as it draws the blueprint, all while making sure the building is stable and looks good.

Here is how Morph works, broken down into simple concepts:

1. The Problem: The "Fixed Room Count" Trap

Most current AI models for designing molecules are like those rigid architects. They assume the number of atoms (the "rooms") is fixed from the start.

The Issue: Many important properties of a molecule (like how well it dissolves in fat or how strong it is) are directly linked to its size. If you can't change the size, you can't easily steer the AI to create a molecule with a specific property.
The Result: If you ask a fixed-size model to design a giant molecule, it might fail because it was only trained on small ones. It's like asking a model that only knows how to build sheds to suddenly build a cathedral; it doesn't know how to scale up.

2. The Solution: "Morphing" Like Clay

Morph treats molecule design like sculpting with clay.

The Process: Instead of just drawing a static picture, Morph starts with a rough shape. As it refines the design, it can:
- Insert: Add new atoms (add a new room).
- Delete: Remove unnecessary atoms (knock down a wall).
- Substitute: Change an atom's type (turn a wooden door into a steel one).
- Move: Shift atoms around to make the structure fit together perfectly.
The Magic: It does all this at the same time. It's not just adding rooms; it's rearranging the whole building while deciding how many rooms it needs.

3. The Secret Sauce: "Unbalanced Optimal Transport"

How does Morph know which atom in the starting sketch corresponds to which atom in the final molecule, especially when the numbers don't match?

The authors use a mathematical trick called Unbalanced Optimal Transport.

The Analogy: Imagine you have a pile of red bricks (your starting sketch) and a pile of blue bricks (your target design). The piles are different sizes.
The Old Way: You try to match every red brick to a blue brick perfectly. If the piles are different sizes, the math breaks.
Morph's Way: It says, "Okay, let's match as many bricks as we can. For the extra blue bricks, we'll just create them out of thin air (insertion). For the extra red bricks, we'll discard them (deletion)."
This allows the AI to smoothly transition from a small, messy sketch to a perfect, complex molecule, even if the starting and ending sizes are totally different.

4. What Morph Can Do (The Results)

The paper tests Morph on two main tasks:

Building from Scratch (De-novo Design): Morph can create valid, stable molecules just as well as the best fixed-size models. It doesn't sacrifice quality for flexibility.
Steering the Design: This is where Morph shines.
- Scenario: You want a molecule that is very "oily" (high lipophilicity). Usually, oily molecules are bigger.
- Fixed Model: If you ask a fixed-size model to make a big oily molecule, and it was only trained on small ones, it might crash or produce garbage.
- Morph: You tell it, "Make me an oily molecule." Morph realizes, "Okay, to be oily, I need to be bigger," and it adds more atoms during the generation process. It successfully creates molecules that are 40% larger than anything it saw in its training data.
Scaffold Decoration: Imagine you have a specific Lego base (a "scaffold") and you want to build a unique structure on top of it. Morph can take that base and intelligently add the right pieces to complete the molecule, whereas other models struggle to modify the size of the structure while keeping the base intact.

5. The Catch (Limitations)

The paper is honest about the downsides:

It's Harder to Train: Because Morph has to learn how to add, remove, and move things all at once, it takes more time and computing power to train (about 2.5 times longer than some other models).
It's Slower to Generate: Because it has to make many small decisions (like adding one atom at a time) rather than just drawing a fixed picture, it takes more steps to finish a single molecule.

Summary

Morph is a new AI tool that breaks the "fixed size" rule in molecular design. Instead of being stuck with a pre-determined number of atoms, it can grow and shrink molecules on the fly. This allows scientists to ask for molecules with specific properties (like size or solubility) and have the AI figure out the right number of atoms to make it work, opening the door to discovering entirely new types of chemicals that previous AI models couldn't reach.

Technical Summary: Generative Molecular Morphing for Flexible-Size Design via Unbalanced Optimal Transport

1. Problem Statement

Current generative models for 3D molecular design, particularly diffusion and flow-based approaches, face a critical algorithmic limitation: they operate on a fixed number of atoms ( $n$ ). This constraint prevents the models from learning the joint distribution of molecular properties and molecular size. Since many chemical properties are intrinsically linked to the number of atoms, fixing the size artificially restricts a model's expressivity and steerability. Consequently, these models struggle to navigate complex relationships between structure and properties, fail to integrate structural priors (like scaffolds) effectively, and cannot generate valid molecules outside the training distribution (OOD) regarding size.

2. Methodology

The authors introduce Morph, a flexible-size generative model for geometric graphs that dynamically adjusts the number of nodes (atoms) during the generation process. The methodology integrates continuous dynamics with discrete jumps, formulated as a jump-flow process.

Core Components:

State Space: The model operates on a state space $\mathcal{G}$ defined as the disjoint union of quotient spaces for geometric graphs of varying sizes, accounting for $SE(3)$ (translation/rotation) and $S_n$ (permutation) symmetries.
Generative Dynamics: The process combines:
1. Continuous Flow: Governed by a velocity field $v_t$ that moves atom positions.
2. Discrete Jumps: Governed by a Continuous-Time Markov Chain (CTMC) with rate density $u_t$ , enabling node insertion, deletion, and type substitution.
Training via Conditional Generator Matching (CGM): To bypass the intractability of the marginal generator, the model uses CGM. It learns to regress conditional velocity and jump rates along paths between paired endpoints $(g_0, g_1)$ , where $g_0$ is a prior graph and $g_1$ is a target molecule.
Unbalanced Optimal Transport (UOT) for Alignment: A key innovation is the use of UOT to align geometric graphs of different sizes ( $n_0 \neq n_1$ $n_{0} \neq = n_{1}$ ) before defining interpolation paths.
- The UOT objective minimizes a cost function involving transport costs, mass deletion penalties ( $w_{del}$ ), and mass insertion penalties ( $w_{ins}$ ).
- Edge edit costs are set to zero, justified by the ability to infer bond topology from 3D structure and valency.
- The resulting transport plan is discretized using the Hungarian algorithm on an augmented space to determine exact edit sets (matches, insertions, deletions).
- Matched positions are rotationally aligned via the Kabsch algorithm to ensure $SE(3)$ invariance.
Interpolation Paths: Based on the UOT matching, the model defines conditional probability paths for three modalities:
- Matched Nodes: Positions follow a Gaussian path; types follow a scheduler.
- Insertion Nodes: Nodes are sampled at specific event times with positions drawn from a Gaussian centered on the target.
- Deletion Nodes: Nodes persist until a specific event time, after which they and their edges are removed.
Loss Function: The training objective minimizes a Bregman divergence between predicted and target conditional generators, comprising terms for movement (MSE), substitution (cross-entropy), insertion (Poisson NLL and GMM), deletion (binary cross-entropy), and charge prediction.

3. Key Contributions

Extension to 3D Geometric Graphs: The work extends flexible-size sequence generation techniques to 3D geometric graphs, enabling the seamless integration of structural priors (e.g., scaffolds) into the generative process.
UOT-Based Matching Algorithm: The authors propose a matching algorithm based on Unbalanced Optimal Transport that aligns geometric graphs of different sizes to define target probability paths for training, handling the trans-dimensional nature of the problem.
Enhanced Steerability and OOD Generation: The paper demonstrates that Morph matches state-of-the-art (SOTA) fixed-size models in unconditional generation while offering superior property steerability. Crucially, it enables successful Out-of-Distribution (OOD) generation in regimes (e.g., larger molecule sizes) where previous models fail.

4. Results

De-novo Design: On the QM9 dataset, Morph achieves competitive metrics (Atom Stability, Molecular Stability, Validity, Uniqueness) compared to SOTA fixed-size models (FlowMol, MiDi, EQGAT-diff, SemlaFlow), often using fewer or equal Number of Function Evaluations (NFE). It also performs well on the more complex GEOM-Drugs dataset.
Property Steering:
- Size Conditioning: When conditioned on atom counts significantly higher than the training distribution (OOD), Morph maintains high validity, generating molecules with up to 40% more atoms than the largest training example. Fixed-size models fail to extrapolate in this regime.
- LogP Conditioning: The model successfully steers the size of generated molecules based on the log partition coefficient (logP), recovering the training distribution's size-marginal and generating valid molecules in low-support or OOD logP regions.
Scaffold Decoration: In a lead optimization task, Morph successfully extends small chemical scaffolds into complete molecules while preserving the core substructure, achieving high validity and uniqueness.

5. Significance and Claims

The paper claims that Morph addresses a fundamental limitation in current generative molecular design: the inability to adapt molecular size during generation. By learning the joint distribution of size and properties, Morph unlocks:

Unparalleled Sampling Flexibility: The ability to generate molecules of varying sizes from a single model.
Improved Discovery Potential: The capacity to explore chemical spaces outside the training distribution (OOD), which is essential for discovering novel molecules with specific, unseen properties.
Structural Prior Integration: The ability to use scaffolds as starting points for flexible expansion, a common requirement in drug discovery.

The authors acknowledge limitations, including increased training costs due to combinatorial complexity, the need for more epochs to reach SOTA performance, and the current reliance on high integration steps due to the CTMC nature of the model. They position Morph as a promising building block for accelerating discovery in chemical sciences, specifically for tasks requiring flexible-size generation.