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

2025-11-12

AI Summary

Existing diffusion and flow-matching methods struggle to model generative tasks with dynamically varying element counts—e.g., variable-length text or protein chains. To address this, we propose Branching Flows, the first flow-matching framework extended to stochastic tree structures that explicitly support splitting (birth) and deletion (death) operations, thereby modeling evolutionary processes over forest-like state spaces. Our approach unifies discrete sets, Euclidean spaces, smooth manifolds, and hybrid modalities without requiring a pre-specified sequence length. By parameterizing dynamics via random binary trees, we tightly couple diffusion mechanics with branching birth–death kinetics, enabling explicit modeling of element insertion and removal. Empirically, Branching Flows demonstrate stable training and strong distributional fidelity on small-molecule, antibody-sequence, and protein-backbone generation tasks—significantly advancing both performance and flexibility for variable-length structural data generation.

Technology Category

• Machine Learning: Structured Learning
• Natural Language Processing: Generation
• Computer Vision: Diffusion Models for Vision

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Social Networks and Social Media: Generative AI / large language models and their impact on social systems
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

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.

Problem

• Modeling variable-length sequences in generative frameworks
• Handling stochastic element creation and deletion during generation
• Extending flow matching to multimodal discrete and continuous spaces

Innovation

• Learns branching and deletion rates on binary trees
• Controls sequence element count during generation process
• Composes with flow matching across multimodal spaces