LeakGFN: Robust Molecular Generation in Generative Flow Networks via Flow Decomposition

Generative Flow Networks (GFlowNets) have emerged as a powerful framework for molecular generation, sampling diverse candidates proportionally to a reward function. However, the vast chemical space necessitates truncating trajectory length, forcing models to treat incomplete molecular fragments as terminal states alongside valid molecules. This conflation distorts the learned distribution by allocating probability mass to chemically meaningless states. We propose LeakGFN, a dual-head architecture that decomposes flow into two components: a chemical head modeling flow over the full chemical space, and a valid head estimating the fraction of flow reaching valid molecules within the truncation boundary. Through this decomposition, the valid head implicitly learns molecular reachability without explicit supervision. We prove that LeakGFN recovers the correct distribution over accessible molecules under mild assumptions. Experiments on five molecular optimization tasks demonstrate consistent improvements over flow matching baselines, achieving state-of-the-art performance on four out of five tasks. Our module integrates as a plug-and-play enhancement into existing frameworks, improving performance on both pocket-conditioned and multi-objective generation tasks.

Designing new drug molecules means searching an astronomically large space of possibilities, far exceeding the number of atoms in the universe. AI systems called generative models can help by learning to propose promising drug candidates. One popular approach builds molecules step by step, assembling a structure from building blocks. However, because the search space is so vast, these systems must impose a limit on how many building steps they can take. This creates a problem: molecules that are still incomplete when the step limit is reached get mixed in with finished ones, confusing the AI about what a good molecule looks like. We developed LeakGFN, a method that teaches the AI to distinguish between complete molecules and unfinished fragments. Our system uses two separate components, one that models the overall construction process and another that estimates whether a molecule can be completed within the allowed number of steps. This separation allows the AI to focus its attention exclusively on valid, complete molecules. In experiments across five drug design tasks, LeakGFN consistently outperformed existing methods, achieving the best results on four tasks. Because it works as a simple add-on, it can be easily combined with other molecular generation systems to improve their performance.