Atomic GFlowNets (A-GFN)

Generative Flow Networks (GFlowNets) learn a stochastic policy that samples molecules in proportion to a reward, which makes them well suited to drug design problems where a diverse set of good candidates matters more than a single optimum. Almost all molecular GFlowNets to date build molecules by joining predefined fragments. That keeps trajectories short and training stable, but it confines the model to whichever pockets of chemical space the fragment vocabulary happens to cover. Atomic GFlowNets (A-GFN) removes that constraint by making individual atoms the action space, at the cost of far longer generation trajectories and a much larger state space that naive training collapses on.

The paper's answer is pretraining. A-GFN is first trained without any task labels against inexpensive, computable molecular descriptors — quantitative estimate of drug-likeness (QED), topological polar surface area (TPSA), synthetic accessibility (SAS), and ring count — used as proxy rewards. The resulting policy is a general-purpose prior over drug-like chemical space that can then be goal-conditioned and finetuned toward a specific downstream property without retraining from scratch. A-GFN was developed by researchers at the University of British Columbia, Mila, Simon Fraser University, and Recursion, and released as a preprint in March 2025 with MIT-licensed code.

#Key Features

• Atom-level action space: Molecules are built one atom at a time from a chemistry-relevant element set (C, N, O, S, P, F, Cl, Br, I plus implicit hydrogen), so generation is not bounded by a fragment library.
• Unsupervised pretraining on proxy rewards: QED, TPSA, SAS, and ring count are cheap to compute for any molecule, which lets the prior be trained on unlabeled compound libraries rather than assay data.
• Hybrid online–offline training: Each minibatch mixes on-policy trajectories with expert trajectories reconstructed from a drug-like compound set, and a maximum-likelihood term pulls samples toward make-on-demand chemical space.
• Goal-conditioned finetuning: The same prior is adapted to property optimization, property targeting under shifted ranges (Dynamic Range Adjustment), and property-constrained optimization.
• Relative Trajectory Balance objective: Finetuning with RTB treats the pretrained policy as a prior and learns a posterior proportional to prior times new reward, which prevents the catastrophic forgetting of drug-likeness seen with a plain Trajectory Balance objective on multi-objective tasks.

#Technical Details

The policy is a graph neural network built from graph attention layers, augmented with additive-aggregation convolutional layers whose output is concatenated into each attention layer's input to support counting-like behaviour over substructures. The main prior was pretrained on ZINC250K — 250,000 commercially available drug-like compounds — for 12 days on four A100-40GB GPUs; an alternative offline-only prior was trained on 6.75 million Enamine compounds for 18 days on the same hardware, and it is this checkpoint, at roughly 9.36 million parameters, that is distributed as AGFN-large. Finetuning runs take 24 hours on a single V100. Against a fragment-based GFlowNet conditioned on identical property ranges, pretrained A-GFN reaches 196 distinct scaffolds per 200 sampled molecules versus 113, with uniqueness of 0.96 and diversity of 0.88. On property targeting, shifting the TPSA window away from the pretraining range gives a 94.5% success rate for the finetuned model against 0% for an A-GFN trained from scratch. In docking-based de novo design against parp1, fa7, 5ht1b, braf, and jak2 with QuickVina2 rewards, finetuning the hybrid prior improves hit ratio over the strongest published baseline by 31.7% to 148.4% depending on target, and novel hit ratio by 34.7% to 188.3%.

#Applications

A-GFN targets the two workhorse tasks of early small-molecule discovery. In de novo design, a chemist specifies a docking target and property window and samples a structurally diverse candidate set rather than a single optimum. In lead optimization, the Bemis-Murcko scaffold of a known lead seeds generation and the model decorates it atom-wise to improve docking score and ligand efficiency while preserving the scaffold. Because the pretrained prior transfers, teams can stand up a new ADME-T objective — Caco-2 permeability, LD50, lipophilicity, solubility, plasma protein binding, or microsomal and hepatocyte clearance — from a small labeled dataset instead of training a generator per endpoint.

#Impact

A-GFN is the clearest demonstration to date that the pretrain-then-adapt recipe transfers to GFlowNets, turning what had been per-task samplers into a reusable generative prior. The model itself is small by foundation-model standards, and every result is in silico: rewards come from docking scores and property predictors, with no wet-lab validation. Complex endpoints remain hard — on toxicity optimization neither the finetuned nor the from-scratch model discovers any high-reward modes, which the authors attribute to the reward scheme rather than the architecture. Finetuning also anchors generation to the prior, so an offline-only prior trades structural novelty for raw hit rate. The preprint has not yet been peer reviewed.