Routing by Reaching: Composition of Pre-trained GFlowNets for Multi-Objective Generation
Abstract
Generative Flow Networks (GFlowNets) learn to sample diverse candidates in proportion to a reward function, making them well-suited for scientific discovery, where exploring multiple promising solutions is crucial. Further extending GFlowNets to multi-objective settings has attracted growing interest as real-world applications often involve multiple, conflicting objectives. However, existing approaches require joint training for each combination of objectives, meaning that any change in the objective set necessitates retraining from scratch. We propose a framework that composes pre-trained GFlowNets at inference time, enabling rapid adaptation without fine-tuning or retraining. Importantly, our framework is flexible, capable of handling diverse reward combinations ranging from linear scalarization to complex nonlinear operators, which are often handled separately in previous literature. We prove that our method exactly recovers the target distribution for linear scalarization, and quantify the approximation quality for nonlinear operators through a distortion factor. Experiments on a synthetic 2D grid and real-world molecule generation tasks demonstrate that our approach achieves performance comparable to baselines. The code is available at https://github.com/ml-postech/gflownet-composition.
Lay Summary
Designing a new drug means juggling several goals at once. A single molecule must stick tightly to its biological target, be easy to manufacture, and be safe enough to act like a real medicine. AI tools can generate promising candidates, but today they must be retrained from scratch whenever researchers change which goals matter, or how much each one counts. We developed a way to combine separate AI models, each an expert on a single goal, instantly and with no retraining. Picture building a molecule piece by piece, like assembling a route one turn at a time. At each step, every expert suggests where to go next, and we give more weight to the expert for whom the molecule built so far is most relevant. By routing the construction toward whichever expert cares most at each moment, the combined system produces molecules that balance all the goals together. This turns single-goal AI models into reusable building blocks that can be mixed and matched on demand. Adding a new goal now means training just one small model instead of redoing everything. This saves time and energy, and makes multi-goal molecular design far more accessible for drug discovery.