FragFlex: Fragment-Based Flexible Molecular Generation for Exploring Vast Chemical Space

AI summary

Overview

Research area: Machine learning for molecular generation — specifically discrete graph diffusion models for de novo drug design.

Technical level: Advanced. The paper assumes familiarity with denoising diffusion models, discrete state-space transition kernels, graph transformers, and cheminformatics concepts such as BRICS fragmentation, valence constraints, and RDKit sanitization.

Scope: FragFlex is a flexible-length discrete diffusion model that inserts and deletes nodes at the fragment level, jointly learning fragment count, fragment identity, and inter-fragment connectivity, then reconstructing an atom-level molecule with a neural attachment-site predictor and a valence-aware assembler.

What This Paper Is About

Drug-like chemical space is estimated at roughly 10^33 to 10^60 molecules, so generative models must propose novel structures efficiently rather than enumerate. Most molecular graph diffusion models fix the number of nodes before denoising begins, so molecular size must be chosen in advance; atom-level insertion/deletion schemes exist but struggle to maintain chemical validity while simultaneously learning size, atom types, and bonds. FragFlex's goal is to combine the chemical inductive bias of fragment-based generation with size-adaptive diffusion, letting molecular size and composition co-evolve within a single generative trajectory.

Key Contributions

1. Fragment-level flexible-size diffusion. FragFlex performs node insertion and deletion at the fragment level rather than the atom level, so the number of fragments changes during denoising while fragment identities and inter-fragment connectivity are denoised jointly.

2. Learned attachment-site prediction replacing lookup tables. Unlike FragDiffusion, which assembles molecules using a fragment-pair attachment-mode lookup table precomputed from training-data attachment frequencies, FragFlex represents fragment connectivity with binary edges and directly predicts the attachment atoms with a neural site head.

3. Valence-aware neural assembler. A top-k, beam-search assembler masks chemically infeasible attachment atoms, renormalizes the site distribution, scores candidate atom pairs with a log-probability sum, and prunes branches that violate RDKit valence checks — filtering learned proposals with chemical constraints at inference time.

4. Large-scale evaluation at 10^7 molecules. The authors generate up to 10 million molecules per model and report validity, uniqueness, QED, SA, distinct valid molecule counts, and Bemis–Murcko scaffold counts, rather than the tens-to-hundreds-of-thousands scale typical of prior evaluations.

Main Findings

• Near-perfect validity sustained across scales: FragFlex holds 99.67% validity at 10^5, 10^6, and 10^7 generated samples. For comparison, DiGress reaches 71.36% and GrIDDD 88.22% at 10^7, while JT-VAE reaches 99.75% and FragDiffusion 96.18%.

• Highest distinct valid molecule count: At 10^7 samples FragFlex yields 9,520,600 distinct valid molecules (|V ∩ U|), the highest of all six models, ahead of FragDiffusion (9,265,827) and JT-VAE (8,916,975).

• Best validity–uniqueness balance at scale: FragFlex retains 95.52% uniqueness at 10^7 samples while keeping 99.67% validity, and is described as the only model above the 95% unique-valid yield contour in the validity–uniqueness plot. JT-VAE, despite the highest validity at every scale, shows the largest uniqueness decline (down to 89.40% at 10^7).

• Best QED and SA at n = 10^5: FragFlex achieves the highest QED (0.764) and lowest SA (3.021) among all evaluated atom- and fragment-based models — better than JT-VAE (0.762 QED, 3.418 SA) and FragDiffusion (0.609 QED, 6.840 SA).

• Scaffold diversity is a trade-off, not a maximum: FragFlex does not produce the most unique Bemis–Murcko scaffolds (67,250 at 10^5, 516,485 at 10^6, 3,492,520 at 10^7), trailing GrIDDD (77,285 / 691,875 / 5,889,181), but attains the best QED/SA profile while still generating millions of unique scaffolds.

• Novelty is not discriminative: Novelty is omitted from the main comparison because it is nearly saturated at 100% across all evaluated models.

• Ablations favor the default configuration: At n = 10^5, default FragFlex gives 99.67% validity, 99.94% uniqueness, and 99,619 distinct valid molecules. Top-1 assembly selection reduces uniqueness (99.22%, 98,888), moving the insertion peak to the midpoint s_p = T/2 slightly raises validity (99.76%) but lowers uniqueness (99.32%, 99,076), and a multi-fragment terminal state maximizes uniqueness (99.97%) at the cost of validity (99.40%, 99,370).

Methodology in Plain English

Representation. Each training molecule is broken into BRICS fragments, keeping the K = 100 most frequent fragment types as the vocabulary. A molecule becomes a coarse graph whose nodes are fragments, whose edges indicate that two fragments are connected, and whose attachment sites record which local atom in each fragment forms the bond. The attachment-mode lookup table from the prior work is used only during preprocessing to derive these atom-level targets; it is not used at inference.

Forward corruption. A root fragment is chosen at random and protected. Every other fragment is assigned a deletion time drawn from a logistic-shaped categorical schedule, so fragments disappear at different points along the diffusion trajectory. Fragments not yet deleted have their identity and edge labels corrupted toward the empirical marginal distribution following a cosine noise schedule. At the final timestep the graph collapses to a single root fragment drawn from the empirical fragment marginal.

Denoising. An edge-aware graph transformer predicts clean fragment identities, binary edges, and directed attachment sites; a separate, smaller count network predicts how many fragments sit at the deletion boundary. Crucially, the boundary placeholders are removed from the count network's input so it cannot simply count them and must infer the required growth from graph structure and timestep. Training minimizes a weighted sum of four cross-entropy terms.

Sampling. Generation starts from one root fragment and runs backward from t = 500 to t = 1. At each step the count network proposes how many placeholder nodes to append, the denoiser predicts their identities and edges, and a categorical posterior is sampled. Because edges are sampled independently, a connectivity repair step links any newly inserted node with no neighbor and, at the end, joins remaining components by adding the highest-scoring cross-component edges.

Assembly. Attachment sites are only needed once, at the last step, so the site head is evaluated only at t = 1 to save computation. The assembler discards atoms that cannot accept a bond, keeps the top-k candidate atoms per fragment edge, scores atom pairs by summed log-probabilities, and runs beam search over edges in order of increasing candidate-set size, discarding any branch that fails RDKit's strict valence check. Up to a fixed number of unique, sanitized, connected products are collected and one is chosen uniformly at random — a deliberate choice that preserves diversity rather than always returning the top-scoring molecule.

Model configuration. T = 500 diffusion steps with a cosine schedule (s = 0.008); logistic deletion/insertion schedule with s_p = 0 and w = 0.05; main denoiser with 7 layers, node width 384, edge width 96, 8 heads; count model with 2 layers, node width 128, edge width 32, 8 heads; maximum graph size N_max = 32 fragments; 200 epochs on 4 GPUs, batch 32 per GPU, dropout 0.05, AdamW with learning rate 2×10^-4 and weight decay 10^-3.

Why This Matters

Impact on research. The paper shows that flexibility and chemical validity need not be traded off against each other: the two capabilities — variable-size diffusion and fragment-based generation — had largely developed separately, and FragFlex bridges them. It also removes a dependency on precomputed attachment lookup tables, replacing a data-derived table with a learned, RDKit-filtered prediction step. The 10^7-scale evaluation gives the field a more demanding benchmark than the usual tens-to-hundreds-of-thousands regime, where uniqueness saturation is rarely exposed.

Real-world applications:

• Generating large virtual screening libraries with high valid-molecule yield for hit finding.
• Scaffold hopping and lead optimization, where fragment composition is the natural design handle.
• Enriching compound collections for under-explored regions of chemical space at scale.
• Providing candidate structures for downstream property prediction and docking pipelines.

Industry relevance. Pharmaceutical and biotech discovery teams care about the fraction of generated molecules that survive sanitization and valence checks, because invalid output wastes downstream compute. A model that returns more than 9.5 million distinct valid molecules from 10^7 samples, with the best QED and SA in the comparison, is directly relevant to library design and to the cost of computational screening campaigns.

Future Directions

• Reinforcement learning with property-based rewards, to balance exploration of new fragment combinations against exploitation of promising molecules.
• Protein-target-conditioned generation, extending FragFlex so fragment composition and molecular size adapt to target-specific binding environments — motivated by the effectiveness of size-adaptive models such as GrIDDD in property targeting.
• Property-guided optimization and downstream drug discovery applications, which the authors list explicitly as future work.
• Closing the scaffold-diversity gap. FragFlex trails GrIDDD and FragDiffusion in unique Bemis–Murcko scaffolds at every reported scale, and the authors frame this as a deliberate balance rather than a solved problem — improving scaffold coverage without sacrificing the QED/SA advantage is an open question. The assembler can also fail even when a valid assembly exists, since both local candidate sets and beam states are truncated.

Target Audience

Machine learning researchers working on discrete diffusion and graph generative models; computational chemists and cheminformatics engineers building de novo design pipelines; and drug-discovery scientists who need to understand the validity/uniqueness trade-offs of generative models at large sampling scales. Readers without background in diffusion processes or molecular graph representations will find the methods section dense, though the results tables are interpretable on their own.

Authors’ abstract

Drug discovery requires efficient exploration of an astronomically large chemical space. Molecular graph diffusion models offer a promising approach to this task, but many require graph size to be fixed before generation. Although atom-level insertion and deletion can enable flexible-size generation, jointly modeling molecular size, atom types, and bonds while maintaining chemical validity remains challenging. We introduce FragFlex, a flexible-length discrete diffusion model that performs node insertion and deletion at the fragment level. During denoising, FragFlex jointly learns the size, fragment identities, and inter-fragment connectivity of the fragment graph. A neural attachment-site predictor and valence-aware assembler then convert the generated fragment graph into an atom-level molecule. In large-scale experiments generating 10 million molecules, FragFlex achieves near-perfect chemical validity while maintaining high uniqueness, demonstrating a strong balance between the two.