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

TL;DR

Branching Flows is a novel generative modeling framework that manages variable-length sequences by evolving elements over a stochastic forest of binary trees, enabling flexible modeling across discrete, continuous, and multimodal spaces.

Contribution

It introduces Branching Flows, which incorporate stochastic branching and dying processes into flow matching, allowing control over sequence length and integration across diverse data domains.

Findings

Effective in small molecule, antibody, and protein generation tasks.

Capable of modeling multimodal, discrete, and continuous data.

Provides a stable learning objective and new generation capabilities.

Abstract

Diffusion and flow matching approaches to generative modeling have shown promise in domains where the state space is continuous, such as image generation or protein folding & design, and discrete, exemplified by diffusion large language models. They offer a natural fit when the number of elements in a state is fixed in advance (e.g. images), but require ad hoc solutions when, for example, the length of a response from a large language model, or the number of amino acids in a protein chain is not known a priori. Here we propose Branching Flows, a generative modeling framework that, like diffusion and flow matching approaches, transports a simple distribution to the data distribution. But in Branching Flows, the elements in the state evolve over a forest of binary trees, branching and dying stochastically with rates that are learned by the model. This allows the model to control, during generation, the number of elements in the sequence. We also show that Branching Flows can compose with any flow matching base process on discrete sets, continuous Euclidean spaces, smooth manifolds, and `multimodal' product spaces that mix these components. We demonstrate this in three domains: small molecule generation (multimodal), antibody sequence generation (discrete), and protein backbone generation (multimodal), and show that Branching Flows is a capable distribution learner with a stable learning objective, and that it enables new capabilities.