Branched transport flows improve deep generative model efficiency
Branched Optimal Transport Amortization
Machine LearningArtificial Intelligence
Summary
Many natural systems like rivers or blood vessels create efficient branching networks. The authors found that common deep generative models don’t capture this branching well, which limits their ability to represent real-world data. They developed a new method that models how data flows can merge and split intelligently, mimicking these natural branching patterns. This helps create better generative models that produce more complex and realistic outputs in areas like biology and images.
What this means in practice
- •For biotechnology developers: Implement branched generative flows to model complex biological data and improve synthetic data generation for biological research.
- •For computer vision engineers: Use branched flow-matching methods to enhance image generation tasks by capturing hierarchical and branching patterns in data.
Authors
Semyon Semenov, Viktor Kovalchuk, Meir Roketlishvili, Albert Baichorov, Fakhri Karray, Martin Takac, Arip Asadulaev
Abstract
Methods of Branched Optimal Transport (BOT) mimic the economy and efficiency of natural tree-like structures, such as those found in rivers and biological systems. These methods are widely applicable for designing efficient networks in society, from river basins and blood vessels to mail and gas distribution systems. However, they remain understudied in the context of designing deep generative models, particularly at a large scale. Standard continuous-time generative models, such as the flow matching approach, fail to capture the inherent hierarchical and branching patterns present in real-world data. Current models provide no mechanism for flows to merge or share pathways to minimize total transport cost. Inspired by the "economy of scale" principle in BOT, we introduce a novel, scalable branched flow-matching algorithm designed to solve the branched optimal transport problem in high dimensions. Our method adapts the Benamou-Brenier continuous-time optimal transport formulation to learn branched generative flows. These flows allow probability mass to aggregate along common pathways before branching out to diverse targets. Parametrized by neural networks, our method effectively learns complex branched generative processes. We demonstrate its effectiveness on challenging high-dimensional tasks in biology and image generation.