Research map / Probabilistic and scientific modeling / Molecules and proteins
Pretraining Generative Flow Networks with Inexpensive Rewards for Molecular Graph Generation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces Atomic GFlowNets (A-GFN), which build molecules atom-by-atom, pretrained on ZINC250K with inexpensive property rewards (QED, TPSA, SAS, rings) via hybrid online-offline trajectory balance, then finetuned with relative trajectory balance for property and docking-based tasks.
Problem
Atom-level GFlowNets for molecular generation have huge action spaces and collapse or produce small, non-drug-like molecules when trained from scratch; fragment-based GFlowNets restrict accessible chemical space.
Approach
Pretrain an atom-based GFlowNet with a mix of offline expert trajectories (backward sampled from drug-like molecules) and online on-policy trajectories, conditioned on cheap molecular properties, plus an MLE regularizer; then finetune with relative trajectory balance to a new reward while keeping the pretrained prior.
Key results
- Pretrained A-GFN found 1252 modes and 196 scaffolds (n=200) versus 0 modes and 113 scaffolds for the fragment-based GFlowNet.
- Finetuning with RTB on property-constrained logP optimization reached 0.9804 top-100 reward versus under 1e-5 for task-trained A-GFN.
- Hybrid-prior finetuning improved hit ratio over the best prior-reported method by 31.7% to 148.4% across five docking targets.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use atom-level A-GFN rather than fragment-based GFN | Molecular property pretraining · Molecular property pretraining | #scaffold (n=200): 113 -> 196 | Helped |
| use atom-level A-GFN rather than fragment-based GFN | Molecular property pretraining · Molecular property pretraining | Div: 0.002 -> 0.88 | Helped |
| finetune pretrained A-GFN rather than train A-GFN from scratch | A-GFN · Molecular weight optimization | Rew (top-100): <10^{-5} -> 0.98 | Helped |
| finetune pretrained A-GFN rather than train A-GFN from scratch | A-GFN · logP optimization | Rew (top-100): <10^{-5} -> 0.98 | Helped |
| finetune pretrained A-GFN rather than train A-GFN from scratch | A-GFN · Toxicity optimization | Rew (top-100): <10^{-5} -> 0.18 | Helped |
| add RTB to pretrained A-GFN finetuning | A-GFN · High TPSA targeting | Success %: 0 -> 94.5 | Helped |
| add RTB to pretrained A-GFN finetuning | A-GFN · Low TPSA targeting | Success %: 0 -> 94.6 | Helped |
| add offline trajectories to pretraining (ablation) | A-GFN pretraining · Pretraining reward | average online reward | Helped |
| use online plus offline trajectories instead of offline-only pretraining (ablation) | A-GFN pretraining · Pretraining reward | average reward | Helped |
| increase offline dataset size from ZINC250K toward Enamine 6.75M (ablation) | A-GFN pretraining · Pretraining reward | final reward | Helped |
| increase temperature β (ablation) | A-GFN pretraining · Molecular property pretraining | convergence speed | Mixed |
| increase the number of Graph Transformer layers (ablation) | A-GFN · Molecular property pretraining | average reward | Helped |
Surprising findings
- Finetuning with plain TB outperformed RTB on simple single-objective optimization tasks, since RTB anchors to the drug-like prior and limits non-drug-like solutions.
- Offline-only priors underperformed in novel hit ratio under strict novelty constraints because RTB biases generation toward known scaffolds.
Limitations
- Catastrophic forgetting of pretrained knowledge under extended finetuning with TB; RTB mitigates but limits exploration.
- Training atom-based models from scratch on docking tasks failed to produce dockable conformers, so docking experiments rely only on finetuning the pretrained prior.
Experimental setup
Datasets: ZINC250K ZINC 1M Enamine 6.75M TDC ADME-T tasks
Benchmarks: TDC ADME-T tasks docking benchmarks for fa7, barf, 5ht1b, parp1, jak2
Metrics: #modes diversity RWTD success percent L1-dist novelty uniqueness hit ratio novel hit ratio top-k reward
Baselines: MolDQN REINVENT 4 Graph MCTS Fragment-based GFlowNet Saturn
Builds on: GFlowNet trajectory balance relative trajectory balance MolDQN outcome-conditioned GFlowNet
Closest papers by meaning
- Generative Flows on Synthetic Pathway for Drug Design (ICLR 2025 · similarity 0.91)
- Generative Flow Networks for Lead Optimization in Drug Design (Student Abstract) (AAAI 2025 · similarity 0.91)
- Navigating Chemical Space with Latent Flows (NeurIPS 2024 · similarity 0.90)
- SynFlowNet: Design of Diverse and Novel Molecules with Synthesis Constraints (ICLR 2025 · similarity 0.90)
- RGFN: Synthesizable Molecular Generation Using GFlowNets (NeurIPS 2024 · similarity 0.89)
- Active Flow Expansion for Out-of-Distribution Discovery: from Theory to Molecules (NeurIPS 2026 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.