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

arXiv:2610.15284 · cs.LG · Submitted 2026-10-11 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "FragFlex: Fragment-Based Flexible Molecular Generation for Exploring Vast Chemical Space".

Jane: The gist: FragFlex introduces a flexible-length discrete diffusion model that performs node insertion and deletion at the fragment level to jointly model molecular size, fragment identities,

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So, we're looking at the paper called "FragFlex: Fragment-Based Flexible Molecular Generation for Exploring Vast Chemical Space". The core thesis is that by moving node insertion and deletion to the fragment level, you can handle variable molecular sizes without having to fix the graph size before generation begins.

Jane: It claims that this approach allows the model to jointly learn the fragment count, their identities, and how they connect throughout the denoising process >

Lu: They introduce FragFlex as a way to bridge a gap between flexible-size diffusion models that work on atoms and fragment diffusion models where you sample the graph size before starting >

Meng: The authors highlight that this design lets molecular size and composition co-evolve within a single generative trajectory, which is pretty ambitious for a model to handle correctly >

Lalam: They also incorporate a neural attachment-site predictor and a valence-aware assembler at the end to convert those generated fragment graphs into actual atom-level molecules, which is key for chemical validity >

Tom: So the main contribution of this work is showing how fragment-based structural flexibility provides a solid foundation for exploring chemical space in ways that previous fixed-size models couldn't manage with the same level of chemical constraint adherence >

Jane: It’s about making sure that when you generate millions of structures, they aren't just random noise but actually follow the rules of chemistry, and this model seems to get that balance right >

Conclusion: Tom: So, we're wrapping up our look at "FragFlex: Fragment-Based Flexible Molecular Generation for Exploring Vast Chemical Space". The authors are showing how they fixed the problem of variable molecular size by making node insertion and deletion happen at the fragment level instead of trying to fix the whole graph size upfront >

Jane: That means they're using a technique where you can add or remove building blocks—fragments—during the generation process itself, rather than just proposing atoms one by one in a fixed-length sequence >

Lu: They combine discrete diffusion with atom-level models to get this flexibility, and the whole goal is to make sure that the model learns the fragment count and how they connect while it's generating, which keeps everything coupled together >

Meng: It’s important to note that they don't just generate random fragments; they use a neural attachment-site predictor and a valence-aware assembler at the end to turn those fragment graphs into actual atom-level molecules, and that assembly step is key for chemical validity >

Lalam: And that assembly part is sophisticated too; they use a scoring mechanism based on PeS probabilities to increase diversity when putting atoms together instead of just picking the highest score every single time >

Tom: The main implication here for drug discovery and materials science is that we can explore much larger chemical spaces efficiently because the generation process itself adapts to the structure it's building >

Jane: If you think about finding new drugs or novel materials, this suggests we can discover structures that might not fit into the traditional fixed-size molds used by older AI models >

Lu: The paper’s contribution is really about showing how fragment-based structural flexibility provides a solid foundation for exploring chemical space in ways that previous methods couldn't manage with the same level of chemical constraint adherence >

Meng: But you gotta remember the practical side; training this model required about ten hours on four NVIDIA L40 GPUs for this fourteen point two million parameter model, so that compute investment is significant >

Sihyun Park, David Ryan Koes

Carnegie Mellon University · University of Pittsburgh

cs.LG

Submitted: 2026-10-11

Updated: 2026-10-11

Comments: 15 pages, 5 figures, 5 tables. Accepted at the NeurIPS 2026 Workshop on AI for Drug Discovery: Bridging the Translation Gap

Code: https://github.com/SihyunPark99/FragFlex

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

The gist: The gist: FragFlex introduces a flexible-length discrete diffusion model that performs node insertion and deletion at the fragment level to jointly model molecular size, fragment identities, and

Key concepts

Discrete Diffusion Model: This is a type of generative model inspired by diffusion processes used in image generation. In this context, it's adapted to work with discrete structures like molecular graphs. It iteratively refines a noisy representation back into a clean molecular structure by predicting the next step in the reverse process.
Fragment-Level Generation: Instead of building molecules atom-by-atom, FragFlex works with reusable substructures called fragments. The model learns to insert and delete these fragments during generation. This approach reduces invalid intermediate structures and allows for flexible molecular sizes by assembling predefined pieces.
Valence-Aware Assembler: This component takes the generated fragment graph and converts it into a final atom-level molecule. It uses a beam search process to score potential connections between fragments, ensuring that the final assembly respects chemical rules like atomic valence while promoting diversity in the resulting molecules.
Joint Learning of Size and Composition: FragFlex is designed so that molecular size (number of nodes) and the identity/connectivity of its fragments are learned simultaneously. This means the model dynamically adjusts how many fragments it uses and which ones it connects during generation, allowing size and composition to co-evolve naturally.

Summary

The gist: FragFlex introduces a flexible-length discrete diffusion model that performs node insertion and deletion at the fragment level to jointly model molecular size, fragment identities, and inter-fragment connectivity while maintaining chemical validity.

Introduction

Drug discovery requires exploring an astronomically large chemical space Estimates of the size of drug-like chemical space vary substantially with the underlying molecular constraints and enumeration strategies ranging from approximately 1033 molecules [25] to 1060 molecules [4]. De novo molecular generation offers a promising strategy for exploring chemical space by directly proposing novel molecular structures Deep generative models can learn molecular distributions from existing compound datasets and generate new molecules Notably, denoising diffusion models have driven substantial progress across image generation [9], discrete text generation [1, 21], and molecular graph generation [33]. Many existing molecular graph diffusion models keep the number of nodes fixed throughout denoising, requiring molecular size to be determined before generation begins Attempts to overcome fixed-length generation are not unique to molecular graphs: masked diffusion models (MDMs) [26] have been extended to support flexible-length text generation through mask-token insertion [15]. For molecular graphs, Ninniri et al. proposed GrIDDD [23], allowing atom nodes to be inserted and deleted during denoising Despite this flexibility, maintaining high chemical validity remains challenging when molecular graph size is allowed to vary through atom-level insertion and deletion This difficulty arises because the model must jointly learn molecular size, atom types, and bond assignments while satisfying strict chemical constraints such as atomic valence. Fragment-based molecular generation offers a natural way to mitigate this challenge Instead of constructing every ring and functional group atom by atom, these methods assemble reusable substructures, reducing invalid intermediates [10, 11]. Fragment-level diffusion and flow models further improve efficiency and scalability by denoising a coarser graph [18]. In particular, Levy Accepted at the NeurIPS 2026 Workshop on AI for Drug Discovery: Bridging the Translation Gap and Rector-Brooks proposed FragDiffusion [19], extending discrete graph diffusion from atoms to fragments. To bridge this gap, we introduce FragFlex, a flexible-length discrete diffusion model that performs node insertion and deletion at the fragment level The model jointly learns fragment count, fragment identities, and inter-fragment connectivity while dynamically adjusting the graph size during denoising. This design allows molecular size and composition to co-evolve within a single generative trajectory while retaining the chemical inductive bias of fragment-based generation. Finally, a neural attachment-site predictor and valence-aware assembler convert the generated fragment graph into an atom-level molecule.

Methods

The overview of FragFlex involves two stages: first, a discrete diffusion model generates a coarse graph by jointly modeling fragment identities, connectivity, and graph size Second, a neural attachment-site predictor and a valence-aware assembler convert this fragment graph into an atom-level molecular graph. The model combines FragDiffusion [19], a discrete diffusion model over fragment graphs, with GrIDDD [23], an atom-level diffusion model that enables variable graph sizes through node insertions and deletions.

The training process involves a forward corruption step where nodes are removed based on a structural schedule The main denoiser fθ predicts clean fragment logits Z X, Z E, and Z S. An auxiliary model gϕ receives the graph input G−t and predicts the number of nodes at the deletion boundary ∆t. The joint objective is optimized using cross-entropy losses for fragment identities, fragment edges, deletion counts, and attachment sites

The reverse sampling process starts from a single root fragment at t = T At each reverse step, gϕ predicts how many DEL∗ placeholders to append The main denoiser fθ outputs clean node and edge logits, which parameterize the reverse distribution pθ(z−1 zt, Gt, t) = X z0 q(zt−1 zt, z0) pθ(z0 Gt, t). Connectivity repair is performed if a newly inserted node has no neighbor in the pre-insertion graph by connecting it to the existing node with the highest predicted probability of a real edge.

Results

In large-scale experiments generating 10 million molecules, FragFlex achieves near-perfect chemical validity while maintaining high uniqueness Table 3 shows that at n = 107, FragFlex maintains 99.67% validity while retaining 95.52% uniqueness. This balance results in 9,520,600 distinct valid molecules, the highest V ∩ U count among all models. FragFlex lies closest to the upper-right corner and is the only model above the 95% unique-valid yield contour

Conclusion

FragFlex introduces a flexible-length molecular diffusion model that performs node insertion and deletion at the fragment level FragFlex jointly predicts the number, identities, and connectivity of molecular fragments A neural attachment-site predictor and valence-aware assembler then convert the generated fragment graph into an atom-level molecule At a scale of 107 generated molecules, FragFlex produces more than 9.5 million distinct valid molecules, demonstrating a strong balance between validity and uniqueness. These results suggest that fragment-level structural flexibility provides an effective foundation for large-scale chemicalspace exploration.

A Detailed Methods

The preprocessing pipeline uses BRICS fragmentation and fragment graph construction FragFlex uses this lookup table only during preprocessing to convert the attachment-mode labels into neural attachment-site targets. Unlike FragDiffusion, FragFlex does not use an attachment-mode lookup table during inference; instead, its neural attachment-site predictor directly proposes the atom-level connection for each generated fragment edge.

The valence-aware assembly uses a beam search process where candidate pairs are scored by lij (a, b) = log PeS ij (a) + log PeS ji(b) This prevents the assembler from always selecting the highest-scoring valid product and increases atom-level assembly diversity for the same fragment graph.

A Ablation Study

The default FragFlex configuration achieves the highest number of distinct valid molecules Selecting only the highest-scoring assembly candidate reduces uniqueness, suggesting that uniform selection helps preserve molecular diversity.

Implementation Details

The model is trained on four NVIDIA L40 GPUs and requires approximately 10 hours of wall-clock time for training the 14.2M-parameter model Sampling was performed on a single NVIDIA GeForce RTX 2080 Ti GPU, requiring approximately 3 hours to generate 105 molecules. The assembly uses k = 4, Bbeam = 512, Pmax = 16.

Fragment Visualization

Figure 4 shows the 100 fragment types used by FragFlex, arranged in descending order of frequency in the ZINC training set.

Generated Molecules

Figure 5 displays 50 random samples generated by FragFlex <ref:2610.

Improvements for AI systems

textbfAbstract flexible-length generation at fragment level by jointly modeling size, identity, and connectivity during denoising makes it possible to generate 10 million molecules with near-perfect chemical validity while maintaining high uniqueness. This enables the AI system to explore vast chemical spaces efficiently without pre-defining molecular size constraints.
textbfConvert fragment graphs into valid molecules using a neural attachment-site predictor and valence-aware assembler, which converts the generated fragment graph into an atom-level molecule, ensuring chemical constraints are satisfied during final assembly. This allows the system to generate chemically valid molecules directly from learned structural representations rather than relying on post-hoc validation or complex constraint satisfaction during generation.
textbfImplement a dynamic graph evolution schedule by jointly modeling fragment identities, connectivity, and graph size, allowing molecular size and composition to co-evolve within a single generative trajectory. This capability permits the system to explore molecular space where the number of fragments changes during generation, which is crucial for discovering novel architectures that might require varying substructure counts.
textbf Utilize a score-based search until P is constructed assembly method, where the final product is sampled uniformly from this pool, instead of always selecting the highest-scoring candidate. This prevents the assembler from always selecting the highest-scoring valid product and increases atom-level assembly diversity for the same fragment graph, ensuring a broader exploration of valid chemical space.
textbf Integrate property-based rewards into reinforcement learning to balance the exploration of new fragment combinations with the exploitation of promising molecules. This suggests an extension where property constraints guide the generation process, allowing for targeted molecular design based on desired physicochemical properties alongside structural novelty.