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

Generative Flow Networks (GFlowNets) sample diverse molecules proportionally to a reward, but the vast chemical space forces trajectory truncation, so incomplete fragments become terminal states alongside valid molecules. We formalize the resulting distortion as flow leakage, the probability mass a converged GFlowNet allocates to these forced terminals. We propose LEAKGFN, which decomposes flow into a chemical head that keeps add-action flow alive at the truncation boundary and a valid head that estimates the fraction of that flow reaching valid molecules from flow matching alone. The decomposition acts on training dynamics near the boundary rather than merely relabeling the target. Separate gradient pathways for flow magnitude and completion decouple the conflicting signals that destabilize single-head models, and our ablations attribute the dominant gain to it rather than to exploration. We also express the residual leakage in closed form and show that at a reward-proportional fixed point the sampler is exact over accessible molecules once conditioned on completion. Experiments on five molecular optimization tasks show the best or tied-best HM on four of them, and the module plugs into existing frameworks, improving pocket-conditioned and multi-objective generation. Code is available at https://github. com/HwanheeKim813/LeakGFN.