Reward-Guided Flow Merging via Implicit Density Operators

ICLR 2026 Conference Submission18862 Authors

Keywords: flow models, diffusion models, model merging, reward-guided fine-tuning

Abstract: Unprecedented progress in large-scale flow and diffusion modeling for scientific discovery recently raised two fundamental challenges: $(i)$ reward-guided adaptation of pre-trained flows, and $(ii)$ integration of multiple models, i.e., model 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. This captures generative optimization tasks requiring information from multiple pre-trained flows, as well as task-aware flow merging (e.g., for maximization of drug-discovery utilities). Our formulation renders possible to express a rich family of implicit operators over generative models densities, including intersection (e.g., to enforce safety), union (e.g., to compose diverse models) and interpolation (e.g., for discovery in data-scarce regions). Moreover, it allows to compute complex logic expressions via generative circuits. Next, we introduce Reward-Guided Flow Merging (RFM), a theory-backed mirror-descent scheme that reduces reward-guided flow merging to a sequential fine-tuning problem that can be tackled via scalable, established methods. 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 on a high-dimensional drug design task generating low-energy molecular conformers.

Primary Area: generative models

Submission Number: 18862