Source-linked AI summary
Branched Optimal Transport Amortization
Semyon Semenov, Viktor Kovalchuk, Meir Roketlishvili, Albert Baichorov, Fakhri Karray, Martin Takac, Arip Asadulaev
TL;DR
Hierarchical and branching data structures are difficult to model with standard continuous-time generative approaches, while BOT has been challenging to optimize at modern scales. BOTA adapts the branched Benamou–Brenier formulation with a soft-atomic relaxation and neural flow matching; the authors report recovering branching geometries and outperforming baselines on complex multimodal distributions.
Problem
BOT has been difficult to scale, while standard continuous-time generative models do not capture hierarchical or branching structures in data.
Source: §1 · pp. 1–2, §5 · p. 5Method
Results
BOTA recovers underlying branching geometries and outperforms baseline methods in reconstructing complex, multimodal distributions.
Source: §9 · p. 9Takeaways & Limitations
The experiments support BOTA as a generative algorithm for learning hierarchical branching processes in biological and image-generation tasks.
Source: §1 · p. 2, §9 · p. 9The differentiable soft-atomic relaxation depends on choices such as proxy atoms and kernel bandwidth that may require task-specific tuning.
Source: §8 · p. 9
Abstract
from arXiv · showhide
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.
1 Introduction
Hierarchical data can require shared transport paths that branch only when needed, whereas common generative flows favor independent straight trajectories. BOTA adapts branched optimal transport to a scalable neural flow-matching solver for hierarchical generative processes.
1 Introduction
- Standard flow-based models encourage straight transport, but such independent motions can poorly fit hierarchical or multimodal data.For taxonomy-related data, computation can be shared along common paths before specializing to particular leaves.Source: §1 · pp. 1–2
- Branched optimal transport uses a subadditive mass cost, often m^α for 0 < α < 1, to favor coalescing network-like paths.This economy-of-scale mechanism is associated with tree-like structures in natural and societal networks.Source: §1 · p. 2
- BOT aligns transport with hierarchical data by rewarding shared prefixes, unlike quadratic-cost OT's independent straight-segment interpolation.Nearby leaves in a taxonomy can travel together for most of the journey and branch only when needed.Source: §1 · p. 2
- Despite its structural appeal, BOT is difficult to optimize at scale: even planar finite-source/sink cases are NP-hard and involve large PDE systems.These computational barriers have kept BOT largely to small or low-dimensional instances.Source: §1 · p. 2
- Modern flow-matching systems operate at large scale, motivating a differentiable, hardware-accelerated BOT solver that respects data hierarchies.Neural amortization is proposed as a way to reuse inferred branched routes across batches and tasks.Source: §1 · p. 2
- BOTA is a continuous, scalable solver for branched optimal transport, parameterized by neural networks using a flow-matching objective.The authors report meaningful hierarchical generative processes on biological and image-generation tasks.Source: §1 · p. 2
2 Background
The background introduces continuous normalizing flows and flow matching as ways to represent and train time-dependent probability paths. It also reviews static and dynamic optimal transport formulations that connect transport costs to paths and velocity fields.
2 Background
- The notation section defines P2(Ω) as probability measures on a compact convex domain with finite second moment and uses ρ for a measure's density when it exists.Vectors are columns, ∥·∥ is the Euclidean norm, ∇· the divergence, and ∥·∥F the Frobenius norm.Source: §2 · p. 2
- A continuous normalizing flow defines a probability-density path through a neural ODE, whose pushforward density satisfies the continuity equation.The continuity equation expresses mass conservation, and integrating the change-of-variables relation permits likelihood calculation.Source: §2 · pp. 2–3
- Flow matching trains a neural velocity field to regress toward a target field generating a desired probability path.Conditional flow matching makes training practical by expressing a marginal path as a mixture of tractable conditional paths.Source: §2 · p. 3
- A Gaussian bridge between coupled samples uses linear interpolation of endpoint locations and a constant drift equal to their displacement.In the stated setup, the coupling is the independent joint distribution of the two endpoints.Source: §2 · p. 3
- Monge–Kantorovich OT minimizes a transport cost, with the Kantorovich relaxation optimizing over couplings when the deterministic-map formulation may be ill-posed.For squared Euclidean cost, the square root of the optimal value is the Wasserstein-2 distance.Source: §2 · p. 3
- Benamou–Brenier OT expresses squared Wasserstein distance as the minimum kinetic energy of a continuous-time flow satisfying the continuity equation.When an optimal map exists, the Wasserstein geodesic is given by interpolation and has constant velocity along trajectories.Source: §2 · p. 3
3 Branched Benamou–Brenier OT
The branched Benamou–Brenier formulation replaces the classical transport cost with a concave mass-dependent cost. Its finite-cost solutions require atomic measures, making direct computation intensive.
3 Branched Benamou–Brenier OT
- The branched formulation uses a concave cost m^αℓ for moving mass m over length ℓ, with α ∈ (0, 1); α = 1 recovers classical OT.The formulation retains the standard endpoint boundary conditions.Source: §3 · pp. 3–4
- Finite branched cost requires the time-dependent measure to be purely atomic for almost every time, with mass concentrated at points.The resulting atomic formulation is computationally intensive, motivating a more efficient relaxation.Source: §3 · p. 4
4 Branched Benamou Brenier Neural OT
The neural BOT approach uses continuous normalizing flows to parameterize transport while relaxing BOT's atomic-mass requirement. A differentiable soft-atomic cost approximates the branched cost using proxy centers and soft assignments.
4 Branched Benamou Brenier Neural OT
- The method uses a continuous normalizing flow to parameterize velocity and obtain a time-continuous density satisfying the endpoint boundary constraint.The branched objective concavifies kinetic action by the density, reflecting the BOT mass-dependent cost.Source: §4 · p. 4
- The learning objective combines a BOT-style dynamic action with a terminal CNF maximum-likelihood fit to the target distribution.The action penalizes kinetic effort while discounting it by a concave density-dependent factor, implementing the move-together-then-split principle.Source: §4 · p. 4
- The atomic constraint complicates direct optimization of the branched Benamou–Brenier objective over continuous-time neural velocity fields.BOT requires an atomic evolving measure, whereas CNFs define continuous density fields.Source: §4 · pp. 4–5
- The soft-atomic relaxation encourages continuous density to cluster around K proxy atoms by softly assigning flow particles to nearby centers.At each time, particles sampled from the flow contribute soft mass and an average velocity for each center.Source: §4 · p. 5
- The resulting instantaneous soft-atomic cost is a differentiable proxy for discretized branched transport and converges to nearest-center hard assignment as σ → 0.The weights are defined by a softmax over squared distances scaled by 2σ².Source: §4 · p. 5
5 Branched Flow Matching via BOT Amortization
BOTA addresses the difficulty of directly optimizing the branched Benamou–Brenier objective by first solving a discrete fixed-batch problem, then amortizing its solution in a continuous-time neural velocity field. The resulting flow-matching model transports new samples using learned branched dynamics.
5 Branched Flow Matching via BOT Amortization
- BOTA first solves a discrete branched-transport problem for a fixed batch, then amortizes its solution into a continuous-time neural network using flow matching.This replaces the straight-line Euclidean transport plan used in standard OT-CFM with a branched plan.Source: §5 · p. 5
5.1 Discrete Branched Transport Solver
- The discrete solver transports source samples to targets over T time steps, minimizing terminal matching error together with time-integrated soft-atomic transport cost.The terminal matching term is a mean-squared error; the solver yields a discrete velocity field and corresponding trajectory for that batch.Source: §5.1 · p. 6
5.2 Amortization via Flow Matching
- To generalize beyond the fixed batch, BOTA interpolates the discrete trajectories into a continuous-time target field and trains a neural velocity field by regression.The learned field can transport new samples from μ0 to μ1 by solving a continuous-time ODE, amortizing the discrete optimization.Source: §5.2 · p. 6
- Figure 2 presents Gaussian-mixture results for increasing numbers of Gaussians, with Flow Matching in the top row and BOTA (α = 0.5) in the bottom row.The caption identifies the methods and layout but does not state a result comparison.Source: Figure 2 · p. 7
6 Experiments
Experiments examine BOTA’s branched transport on Gaussian mixtures, known Tedsim differentiation dynamics, and image-generation tasks. The image experiments include MNIST class generation and FFHQ latent-space interpolation.
6.1 Gaussian Mixtures
- As the number of target clusters grows from 3 to 6 to 18, BOTA forms shared trunks that split into structured, tree-like transport, unlike FM’s crowded, crossing fan-out paths.The comparison uses same-size models and identical training settings.Source: §6.1 · p. 7
6.2 Biological data
- Tedsim provides known branching dynamics; the experiment transports from one progenitor to two terminal states to assess trajectory reconstruction.Figure 3 reports BOTA results on complicated Tedsim at α = 0.5.Source: §6.2 · p. 8, Figure 3 · p. 8
6.3 Image data
- On MNIST, BOTA reveals branching structure when generating classes 0 and 1, and the experiments also map the full dataset from a Gaussian distribution.The hyperparameters match those used in the biology experiment.Source: §6.3 · p. 8, Figure 4 · p. 8
- On FFHQ, BOTA learns a latent-space female-to-male interpolation intended to be smooth and hierarchical, with a pretrained, frozen ALAE used only to decode images during evaluation.The data are split into 60,000 training and 10,000 test samples.Source: §6.3 · p. 8, Figure 5 · p. 9
7 Related Work
Prior BOT formulations offer insight into branched transport but remain computationally intractable for high-dimensional data. Branched generative flows encourage branching, but do not explicitly solve the BOT optimization problem.
7 Related Work
- BOT formulations include traffic plans, irrigation flows, and branched networks, but existing approaches remain computationally intractable for high-dimensional data.The passage describes fiber-based modeling of particle locations over time.Source: §7 · p. 9
- Y-shaped and branched Generative Flows encourage branching, yet do not guarantee a globally optimal branched transport plan for the data.They learn transport hierarchies heuristically during single-stage network training rather than computing or supervising with ground-truth BOT.Source: §7 · p. 9
- Unlike heuristic branched flows, the cited approaches do not compute or supervise with the ground-truth BOT plan.This distinction concerns whether the BOT optimization problem is explicitly solved.Source: §7 · p. 9
8 Limitations and Broader Impact
The method’s soft-atomic relaxation and two-stage optimization introduce tuning requirements and design complexity. The broader-impact discussion also notes potential misuse and dataset-bias risks.
8 Limitations and Broader Impact
- Soft-atomic relaxation geometry depends on proxy atoms and kernel bandwidth, which may require task-specific tuning.Source: §8 · p. 9
- The two-stage discrete optimization and neural amortization procedure adds design complexity.Source: §8 · p. 9
- Potential risks include misuse for synthetic data generation and amplification of dataset biases, especially in sensitive domains.Source: §8 · p. 9
9 Conclusion
BOTA adapts the Benamou–Brenier formulation with a concave economy-of-scale cost and a differentiable soft-atomic relaxation, then amortizes branched solutions through flow matching. Experiments show recovery of branching geometries and improved reconstruction of complex, multimodal distributions against baselines.
9 Conclusion
- BOTA is a scalable generative algorithm aimed at capturing hierarchical, tree-like data structures that standard continuous-time models cannot capture.Source: §9 · p. 9
- A concave economy-of-scale cost encourages probability mass to aggregate along shared trunks before branching toward targets.Source: §9 · p. 9
- A differentiable soft-atomic relaxation makes the approach computationally feasible, and flow matching amortizes the solution into a continuous neural vector field.Source: §9 · p. 9
- Experiments show that BOTA recovers underlying branching geometries and outperforms baselines in reconstructing complex, multimodal distributions.Source: §9 · p. 9
Extended Analysis
The analysis shows that BOTA’s soft cost is differentiable and converges to the hard-assignment branched transport cost as the kernel width shrinks. Experiments illustrate how branching strength shapes transport geometry and report results on image and biological data, alongside training-cost comparisons.
Extended Analysis
- The soft cost is differentiable for σ > 0 and converges to the hard nearest-center assignment cost as σ → 0.When centers coincide with the atom positions and K = B, the expression recovers the atomic branched transport cost; the smooth reformulation supports gradient-based learning of branched velocity fields.Source: Extended Analysis · pp. 11–12
Gaussian Data
- Smaller α encourages shared pathways and pronounced trunk-and-branch structures, while larger α yields increasingly independent trajectories resembling classical optimal transport.The Gaussian ablations examine increasing numbers of Gaussians and different α values.Source: Gaussian Data · p. 12, Figure 6 · p. 13
Image Data
- MNIST samples show clean, well-formed digits across classes, with consistent stroke thickness and minimal visual artifacts.The results are described as maintaining stable trajectories on discrete, multidimensional data and producing reliable class-conditional generations.Source: Figure 7 · p. 12, Image Data · p. 12
- On FFHQ, BOTA achieves lower FID than Flow Matching, with the improvement consistent across seeds.The passage attributes the lower FID to closer alignment of the learned transport field with the underlying data distribution.Source: Table 2 · p. 12, Image Data · p. 12
Biological Data
- BOTA reconstructs Tedsim’s shared progenitor trunk and lineage splits, with substantially lower soft-atomic cost and competitive W1, W2, and RBF-MMD metrics.The model produces smooth trajectories aligned with developmental paths, without mode collapse or spurious branching.Source: Biological Data · pp. 13–14, Figure 8 · p. 14
- BOTA’s training takes roughly 2× the time of Flow Matching, paid once offline, while sampling cost is identical.The discrete solver stage costs O(T · N · K · D), linear in particle count N, with no N^2 term.Source: Biological Data · p. 15, Table 3 · p. 15