📊 statistics

Branching Flows: Discrete, Continuous, and Manifold Flow Matching with Splits and Deletions

Branching Flows is a new generative modeling framework that enables the generation of sequences with variable lengths by allowing elements to stochastically branch or die within a forest of binary trees, while remaining compatible with various flow matching processes across discrete, continuous, and manifold state spaces.

Original authors: Lukas Billera, Hedwig Nora Nordlinder, Jack Collier Ryder, Anton Oresten, Aron Stålmarck, Theodor Mosetti Björk, Ben Murrell

Published 2026-04-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Lukas Billera, Hedwig Nora Nordlinder, Jack Collier Ryder, Anton Oresten, Aron Stålmarck, Theodor Mosetti Björk, Ben Murrell

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 teach a robot how to build something, but there’s a catch: you don’t know how big the final object is going to be.

If you were teaching it to draw a picture, you’d give it a fixed-size canvas. If you were teaching it to write a sentence, you’d tell it to stop after a certain number of words. But what if you wanted it to design a protein molecule, a new drug, or a complex biological chain? In nature, these things don't come in "standard sizes." One protein might have 50 links, and another might have 500.

Current AI models are like artists working on fixed-size canvases. They are great at filling in a specific space, but they struggle when the "canvas" itself needs to grow or shrink.

This paper introduces "Branching Flows," a new way for AI to generate things that change in length.

The Analogy: The Magic Growing Tree

To understand how Branching Flows works, stop thinking about a flat canvas and start thinking about a Magic Tree.

1. The Seed (The Starting Point):
Most AI models start with a "blank" version of the final product. Branching Flows starts with a "seed"—a very simple, small starting point (like a single trunk).

2. The Branching (Growing the Length):
As time passes (during the AI's "thinking" process), the tree begins to grow. A single branch might suddenly split into two. This is how the AI decides, "Actually, I need more pieces here to make this molecule work." It’s not just adding pieces; it’s growing the structure from the inside out.

3. The Pruning (Deleting the Mistakes):
Sometimes, the tree grows a branch that doesn't belong—a "dead end" that doesn't lead to a useful part of the molecule. Branching Flows has a "pruning" mechanism. It can look at a branch and say, "This is a mistake; let's delete it," effectively shrinking the structure back down to the perfect size.

4. The Anchors (Guiding the Growth):
To make sure the tree doesn't grow into a chaotic mess, the AI uses "anchors." Think of these as invisible guideposts in the air. As a branch grows, it is constantly being pulled toward these guideposts, ensuring that even as the tree splits and grows, it is always moving toward the shape of a real, functional molecule.

Why is this a big deal? (The "Infix" Problem)

The researchers highlight a specific problem called "Infix Sampling."

Imagine you are building a bridge. You have the left side finished, and you have the right side finished. Now, you need to build the middle section to connect them. But you don't know how long that middle section needs to be to make the bridge stable!

Old AI models would struggle because they'd try to guess a length first, and if they guessed wrong, the bridge wouldn't fit. Branching Flows solves this. It can start with the two ends and let the "tree" grow the middle section naturally until the two sides are perfectly connected.

Real-World Applications

The researchers tested this "Magic Tree" approach on three very difficult tasks:

Small Molecules (Chemistry): They taught the AI to build tiny molecules. The AI successfully learned how many atoms to add and where to put them to create realistic chemical structures.
Antibodies (Medicine): Antibodies are the "soldiers" of our immune system. They come in many different lengths. The AI learned to grow antibody sequences that look and act like the real ones found in nature.
Proteins (Biology): This is the "final boss" of biology. Proteins are massive, complex 3D shapes. The AI was able to design protein structures that were so realistic that when they were run through other advanced biological simulators, they looked like real, functional proteins.

Summary

In short, Branching Flows moves AI from being a "painter" (stuck on one canvas size) to being a "gardener." It allows AI to grow, split, and prune structures, making it a powerful new tool for designing the complex, variable-sized building blocks of life.

Technical Summary: Branching Flows: Discrete, Continuous, and Manifold

Branching Flows is a novel generative modeling framework designed to address the "variable-length" problem in diffusion and flow matching models. While existing generative models excel at fixed-length data (like images), they struggle with sequences where the number of elements is unknown a priori, such as protein chains, molecular graphs, or language model responses.

1. The Problem: The Unknown-Length Challenge

Current generative paradigms face a dichotomy when handling variable-length sequences:

Autoregressive (AR) Models: Handle length variation well (via stop tokens) but are inherently sequential (left-to-right), making them difficult to scale to complex, non-linear, or continuous multi-modal structures.
Diffusion/Flow Matching Models: Excel at structured, multi-modal, and continuous data but typically require a fixed state dimension.
The "Infix Sampling" Problem: A specific gap exists where a segment of unknown length must be inserted into a sequence conditioned on known flanking regions (e.g., designing an antibody loop between two fixed framework regions). Current "edit-based" or "jump-diffusion" solutions often rely on approximations that do not scale well to complex, multi-modal manifolds.

2. Methodology: The Branching Flows Framework

The authors propose a framework based on Generator Matching, where a conditional process transports a simple distribution ( $x_0$ ) to a data distribution ( $x_1$ ).

Core Mechanism: The Forest of Binary Trees
Instead of a fixed-length vector, the state evolves over a forest of binary trees.

Branching (Splits): Elements can duplicate in-place at bifurcation points. This allows a single root element at $t=0$ to evolve into multiple leaf elements at $t=1$ .
Deletions: Elements can "die" stochastically, allowing the model to prune unnecessary elements.
Stochastic Rates: The model learns the instantaneous rates for both splitting and deletion.

Mathematical Formulation:

Conditional Paths: The model uses a "conditional bridge" construction. Given a pre-sampled tree structure ( $T$ ) and "anchors" ( $A$ )—which act as intermediate targets for the flow—the elements evolve independently along branches.
Anchors: To ensure the flow is stable, internal nodes in the tree are assigned "anchors" (e.g., the centroid of descendant elements). This steers the continuous or discrete elements toward their eventual targets.
Loss Function: The model is trained using a Conditional Branching Flows (CBF) loss, which is a sum of three Bregman divergences:
1. Split Loss: Matches the predicted split intensity to the actual number of splits required by the tree.
2. Deletion Loss: A cross-entropy loss matching the probability of an element being deleted.
3. Base Loss: A standard flow-matching loss for the continuous/discrete evolution of the elements themselves along the branches.

Versatility: The framework is "composable," meaning the base process can be applied to discrete tokens, Euclidean spaces, Riemannian manifolds (like $SO(3)$ for rotations), or multimodal product spaces.

3. Key Contributions

Unified Framework: A single mathematical approach that handles discrete, continuous, and manifold-valued elements simultaneously.
Variable Cardinality: Naturally handles changes in the number of elements without requiring fixed-length padding or ad-hoc "virtual" tokens.
Infix Sampling Capability: Provides a principled way to perform "unknown-length infix sampling," a holy grail in protein and molecular design.
Efficient Sampling: Uses a recursive waiting-time sampling method (based on hazard rates) that allows for efficient inference despite the branching complexity.

4. Results and Empirical Validation

The authors demonstrate the model's efficacy across three complex domains:

Small Molecules (QM9): The model jointly generates atom positions (Euclidean) and atom types (discrete). It outperformed "Transdimensional Jump Diffusion" models in matching molecular properties (e.g., LogP, molecular weight) and atom counts.
Antibody Sequences: The model successfully learned the distribution of amino acid sequences and lengths. It matched the performance of "Oracle-length" models (which are given the correct length upfront) and outperformed "Edit Flows" in sequence content diversity and nativeness.
Protein Structure (BranchChain & BranchSegment):
- BranchChain could generate entire protein chains of varying lengths, preserving symmetry (homodimers) and producing geometrically plausible backbones.
- BranchSegment demonstrated the ability to perform infix sampling, designing a variable-length CDR3 loop in a nanobody while keeping the rest of the structure fixed, successfully avoiding atomic clashes with a binding partner.

5. Significance

Branching Flows represents a significant step toward general-purpose generative modeling for the physical sciences. By moving away from fixed-dimension tensors and toward stochastic tree-based structures, the authors have provided a tool capable of modeling the inherent "growth" and "pruning" processes seen in biology and chemistry. This has profound implications for de novo protein design, drug discovery, and any domain where the complexity of the object is not known until it is fully formed.