A Unified Density Operator View of Flow Control and Merging

Recent progress in large-scale flow and diffusion models raised two fundamental algorithmic challenges: $(i)$ control-based reward adaptation of pre-trained flows, and $(ii)$ integration of multiple models, i.e., flow merging. While current approaches address them separately, we introduce a unifying probability-space framework that subsumes both as limit cases, and enables *reward-guided flow merging*, allowing principled, task-aware combination of multiple pre-trained flows (e.g., merging priors while maximizing drug-discovery utilities). Our formulation renders possible to express a rich family of *operators over generative models densities*, including intersection (e.g., to enforce safety), union (e.g., to compose diverse models), interpolation (e.g., for discovery), their reward-guided counterparts, as well as complex logical expressions via *generative circuits*. Next, we introduce Reward-Guided Flow Merging (RFM), a mirror-descent scheme that reduces reward-guided flow merging to a sequence of standard fine-tuning problems. Then, we provide first-of-their-kind theoretical guarantees for reward-guided and *pure* flow merging via RFM. Ultimately, we showcase the capabilities of the proposed method on illustrative settings providing visually interpretable insights, and apply our method to high-dimensional de-novo molecular design and low-energy conformer generation.

Modern generative AI can propose new molecules and other complex designs, but each trained model usually captures only part of what scientists want. One model may generate chemically plausible molecules, another may favor low-energy structures, and a third may encode safety or design constraints. Today, combining such models while also steering them toward a desired scientific goal is difficult. We introduce a framework that treats two tasks—adapting one model to a goal and merging several models—as the same underlying problem. Instead of averaging the models’ internal parameters, our method combines them by directly shaping the probabilities of what they generate. This lets us define intuitive operations such as “AND” (generate what several models agree on), “OR” (cover what any model can generate), and goal-directed versions that also prefer molecules or designs with better desired properties. We then design Reward-Guided Flow Merging, an algorithm that implements these operations through repeated fine-tuning steps, and prove when this process converges. In experiments, the method behaves as intended in simple visual examples and works on molecular design problems, suggesting a principled way to reuse specialized generative models for scientific discovery.