FragFM: Hierarchical Framework for Efficient Molecule Generation via Fragment-Level Discrete Flow Matching

[Common response to all reviewers]

We thank reviewers for their feedback for their helpful suggestions on clarity and presentation. We will revise the last part of the introduction to include the core contributions of the paper. For clarity, we reemphasize our core contributions below:

• The first general fragment-level flow matching. We embed chemically meaningful substructures into the flow-matching framework for molecular‐graph generation, enabling efficient sampling of complex molecules and stronger property control.
  • Similar or superior performance compared to the previous SOTA models in general molecular generative benchmarks
  • Efficient sampling at very few steps, maintaining high generation quality with 10 steps (≈95% validity, FCD < 1.0 on MOSES), outperforming atom-level baselines even at 500 steps (Fig. 2, Tab. 8).
  • Better conditional generation: With classifier guidance and fragment-bag reweighting, FragFM shows better conditioning ability across every simple and complex properties (Figs. 3-4, App. E.5), while preserving validity under strong conditioning.
• Resolving the key challenges in a fragment-based generative model. To make the approach practical we must (i) avoid the prohibitive cost of enumerating a very large fragment library, (ii) generalize to fragments that never appear in training, and (iii) reconstruct atom-level graphs faithfully.
  • Methodological novelty. We tackle these challenges with (i) a stochastic bag strategy that provides an unbiased, low-cost surrogate for the full library, (ii) a GNN-based fragment encoder that extrapolates to unseen fragments, and (iii) a coarse-to-fine auto-encoder that converts fragment graphs back to atom-level structures with high fidelity.

IMPORTANT NOTE: Because of character limits, we list only the citations that were added during the review. New references are numbered consecutively after the final citation of the main manuscript, while any source already cited in the paper retains its original reference number.

We sincerely appreciate the reviewer’s time and careful evaluation of our work. The reviewer’s thoughtful questions and concrete suggestions on bag-size ablations and autoencoder details were invaluable and led to substantial improvements in the manuscript.

[W1,Q1] Detailed analysis regarding the fragment bag size

Before we introduce ablation study, we briefly summarize the theoretical background of importance sampling and InfoNCE framework. During training, negative candidates are drawn from the unconditional fragment frequency $p(x)$, whereas the positive fragment follows the conditional density $q(x)=p(x|X_t)$. Let

$w(x) = \\frac{q(x)}{p(x)},\\quad \\hat{\\mu}\_N = \\frac{\\sum\_{i=1}^N w(x_i)f(x_i)}{\\sum_{i=1}^N w(x_i)}, \\quad x_i\\sim p,$

where $f$ is a testing function. $\\hat{\\mu}\_{N}$ is a self-normalized importance-sampling (SNIS) estimator of $\\mu=\\mathbb{E}\_{x\\sim q}[f(x)]$.

Standard results (see Chapter 9 of [84]) give

$\\mathbb{E}[\\hat{\\mu}_N]=\\mu + \\mathcal{O}(N^{-1}), \\quad \\mathrm{Var}[\\hat{\\mu}] = \\sigma^{2}/N + \\mathcal{O}(N^{-2}).$

In our setting, we note that $N$ refers to the size of the sampling fragment bag, or $|B\\vert$. Hence inference performance improves monotonically as the candidate-set size $N$ grows, converging to the exact value in the limit (See “Effect of the inference-time bag size” below).

InfoNCE uses the same bag construction with N_\\mathrm{train} in training phase. A larger N_\\mathrm{train} tightens the mutual-information lower bound, but the optimal network simply estimates the density ratio $w(x)$, which is independent of $N_{\\mathrm{train}}$ [44]. In practice, N_\\mathrm{train} controls gradient variance.

In the previous experiments, we set N_\\mathrm{train}=N=384 , guided by preliminary sweeps and theory as a training and inference-time cost–accuracy trade-off (see Reviewer A3wR Q2). Thanks to the insightful review, we additionally ran extensive experiments varying training- and inference-time bag-sizes to systematically analyze the effects of both of them; the results are displayed below. We will add the results and theoretical analysis in the manuscript.

For ablation about the bag sizes, we run three experiments:

1. changing $N$ with fixed N_\\mathrm{train} to study inference-time bag size effects
2. changing N_\\mathrm{train} with fixed $N$ to study training-time bag size effects
3. changing both N_\\mathrm{train}=N to study overall effects of fragment bag sizes (please refer to Reviewer A3wR Q2)

Effect of the inference-time bag size. Back to the reviewer’s original question, we re-ran the MOSES benchmark with the same trained model ($N_{\\text{train}}=384$) while varying the inference bag size $N$ to verify that training-inference bag sizes can differ safely. As predicted, all metrics stabilize once N \\ge N_\\mathrm{train} and degrade gracefully for smaller bags. These results well aligns to the theoretical analysis.

(To isolate the effect of bag size, these generations exclude the detailed-balance term. $N$ is $\vert \mathcal{B} \vert$ during generation, $N_\mathrm{train}$ is $\vert \mathcal{B} \vert$ during training.)

Effect of the training-time bag size. To isolate training-time effects, we additionally trained independent models with different N_\\mathrm{train} on the MOSES benchmark with sufficiently large $N$. During inference, we fixed $N=1536$ for all model. The results shows stable across N_\\mathrm{train}: validity >99.5%, FCD = 0.58-0.65, and Filters ~ 98.5%, suggesting that all trained models with $|\\mathcal{B}|$ larger than 16 can estimate the density ratio $w(x)=q(x)/p(x)$ robust on the $N_{\\mathrm{train}}$.

(To isolate the effect of bag size, these generations exclude the detailed-balance term. $N_\mathrm{train}$ is $\vert \mathcal{B} \vert$ during training.)

[84] Owen, Art B. “Monte Carlo theory, methods and examples.” (2013): 19-22.

[W2] More information regarding coarse-to-fine autoencoder and fragment library

For the implementation details about coarse-to-fine autoencoder, full size of the fragment library, and detailed fragmentization strategy please refer to our reply for Reviewer A3wR’s [W1,Q1]. As many reviewers pointed out limited explanation for coarse-to-fine autoencoder in the main text, we will revise the manuscript to include more information about the coarse-to-fine autoencoder.

[Q2] Why is the curve of QED=0.8 deviates from curves of other target values in Figure 3?

The deviation stems from the dataset prior. On MOSES, the QED distribution is peaked around 0.8 (Figure 9c). Targeting therefore requires a much smaller shift from the training data, so classifier guidance attains lower MAE at lower FCD when targeted to such value. These trends are similarily shown in logP and number of rings conditioning (Figure 12 and 13).

We would like to thank the reviewer for their thoughtful and constructive review with their perspective to significantly strengthen the clarity and impact of our work.

We include a common response to all reviewers at the beginning of our response to Reviewer weUv, re-emphasizing our core contributions.

[W1,Q1] Reading this paper is not easy

We sincerely thank the reviewer for their constructive feedback on improving the clarity and reproducibility of our work. We agree that some implementation details were missing from the main text. To address this, we have made substantial revisions, which we will include in the camera-ready version.

We outline the specific changes below, corresponding directly to the reviewer's points.

1. Implementation Details and Hyperparameters

We apologize for omitting the fragment bag size and fragment library sizes. We will explicitly state all key hyperparameters.

• Fragment Bag Size ($\\vert\\mathcal{B}\\vert$): We will clarify in the text that the fragment bag size for sampling is set to 384.

• Fragment Library Size ($\\vert \\mathcal{F} \\vert$): We will include the following table about fragment library size. Note that full utilization of the extensive training / test fragment libraries is not explored yet in the previous fragment-based approaches used only frequently occuring 100 to 200 fragments.

  	MOSES	GuacaMol	NPGen
  Total	44,083	223,186	133,823
  Training	27,760	145,971	84,731
  Valid	2,248	7,473	3,709
  Test	2,541	22,508	11,712

• Fragmentation and Bond Types: We will state clearly that our fragmentation strategy is based on BRICS rules [38] and, importantly, that our fragment library and inter-fragment connections exclusively use single bonds. This design choice simplifies the generative process.

2. Methodological Clarity: Coarse-to-Fine Autoencoder

We agree with the reviewer that the original description of our autoencoder was insufficient. We have rewritten this section to be more self-contained and clear. The revised text below now provides a detailed information of:

1. The model's training strategy.
2. The specific fragmentation rules and resulting library.
3. The decoding process, explaining how the latent vector, fragment-level graph and the Blossom algorithm are integrated to reconstruct the final atom-level graph.

Revised version of 3.2 coarse-to-fine autoencoder

While a fragment-level graph ($\\mathcal{G}$) offers a higher-level abstraction of molecular structures, it also introduces ambiguity in reconstructing atomic connections. Specifically, a single fragment-level connectivity ($\\mathcal{E}$) can map to multiple distinct, valid atom-level configurations. To achieve accurate end-to-end molecular generation, it is therefore crucial to preserve atom-level connectivity ($\\mathbf{E}$) when forming the fragment-level representation. Drawing on hierarchical generative frameworks [40,41,42], we employ a coarse-to-fine autoencoder. The encoder first decomposes an atom-level graph ($\\mathbf{G}$) into a set of predefined fragments using BRICS rules, where we only considered single bond-based fragmentation rule. It then compresses the precise inter-fragment atomic connectivity information, which are lost during fragmentation, into a single continuous latent vector, $z$. Consequently, the decoder's task is to reconstruct the original molecule by predicting the atom-level edges between fragments, guided by both $\\mathcal{G}$ and $z$.

We train our autoencoder using a combined loss function designed for high-fidelity reconstruction and a well-behaved latent space. The primary objective is a reconstruction loss that ensures the predicted atomic links between fragments faithfully restore the original molecular structure. We supplement this with a minimal Kullback-Leibler (KL) divergence penalty on the latent variable. Finally, to translate the decoder's output into a chemically valid molecular graph, we employ the Blossom algorithm [43]. In detail, the decoder predicts a matrix of probabilities for all potential atomic connections between the "junction atoms" of adjacent fragments. We interpret this as a maximum weighted matching problem, where the edge weights are the log-probabilities predicted by our model. The Blossom algorithm then efficiently determines the optimal set of one-to-one atomic connections that maximizes the joint probability, guaranteeing a valid and accurate reconstruction of the final molecule.

These comprehensive updates will be integrated into the final manuscript. We are confident that these revisions directly address the reviewers’ concerns and make our work significantly more accessible and understandable to the NeurIPS readers.

[W2,Q3] Mentioning more related works and new baselines

We thank the reviewer for their constructive feedback and for highlighting these relevant works. A key distinction is that the suggested papers focus on fragment-based language models, whereas our work proposes a fragment-based graph model. This fundamental difference in data modality (sequential strings vs. explicit graphs) guided our initial set of baslines.

In order to conduct a fair additional benchmarking analysis in terms of training data, we re-trained the models on the MOSES benchmark. Here, GenMol was excluded from the comparison because they did not provide the training code. We trained Sc2Mol for 25 epochs with the provided default settings and trained SAFE-GPT-20M with 20 million parameters, following the detailed hyperparameter settings in Section 4.1.5 of the original paper [83]. Additionally, the validity correction module of Sc2Mol, which explicitly fixes the validity of generated SMILES, is not publicly available. Therefore, we only report the results of Sc2Mol without validity correction. After generating 25,000 molecules for each model with three different seeds, we applied the same post-processing strategy to the molecules.

The results indicate that these methods have limitations in capturing the target data distribution. Except for novelty, for which we provide additional insight in Q1 of the EQJM reviewer's comments, FragFM outperformed all other fragment-based language models.

Given these results, we believe our method offers a distinct and effective alternative in the landscape of fragment-based generation.

Furthermore, we will add a discussion of these methods to our related work section.

Related Work

Recently, several fragment-based molecular generation models that use sequence-based representations have also been developed. For instance, SAFE-GPT [83] uses a transformer architecture on the SAFE (Sequential Attachment-based Fragment Embedding) representation, which linearizes molecular graphs into sequences of fragments. Similarly, GenMol [84] builds upon the SAFE representation but employs a discrete diffusion model for generation. Another approach, Sc2Mol [85], utilizes a two-stage process where a variational autoencoder first generates a carbon scaffold, which is then decorated with specific atoms and bonds by a transformer model. While these methods have shown success in generating valid molecules, they are conceptually different from our approach, which operates directly on graph representations rather than linear sequences.

[83] Noutahi, Emmanuel, et al. "Gotta be SAFE: a new framework for molecular design." Digital Discovery 3.4 (2024): 796-804.

[84] Lee, Seul, et al. "Genmol: A drug discovery generalist with discrete diffusion." arXiv preprint arXiv:2501.06158 (2025).

[85] Liao, Zhirui, et al. "Sc2Mol: a scaffold-based two-step molecule generator with variational autoencoder and transformer." Bioinformatics 39.1 (2023): btac814.

[Q2] Fragment bag size-performance evaluation

Because the fragment-bag size $|\\mathcal{B}|$ governs our Monte Carlo approximation of the full-fragment library, increasing the fragment-bag size would lower the variance (improving generation fidelity) at the expense of compute. We therefore ran a ablation study on MOSES with $|\\mathcal{B}|=16,32,64,128,256,512$, and will include it in the manuscript. To isolate the effect of bag size, these generations exclude the detailed-balance term. As $\vert \mathcal{B} \vert$ increases, the metrics improve—higher validity/Filters/SNN and lower FCD—and converges, as expected. We also provide theoretical analysis along with additional results on varying trian / test bag size on Reviewer weUv’s [W1,Q1].

I thank the authors for taking the time to write a comprehensive response addressing all the concerns. I appreciate the extra effort authors have put into comparing with related models. Given that the paper would be updated according to the responses, I would like to increase my review score.

Question: In the table comparing performance of related work the model Sc2Mol (w/o VC) has generated molecules with a novelty score of 98.6 which is higher than SAFE-GPT and FragFM. However, novelty 90.9 of SAFE-GPT is made bold. Is there any reason for Sc2Mol to have such a high novelty score? What is the percentage of novel and valid moelcules?

Dear Reviewer A3wR,

We thank the reviewer for recognizing our new experiments. Your feedback helped us broaden the related-work discussion and make the paper more persuasive. We also appreciate the initial comments on fragment-bag size, an important aspect of our method, thereby provided additional analysis and greatly improved the manuscript’s quality.

We apologize to mis-represent the result. As we cannot fix the initial rebuttal now, we fix the benchmark table and revised it below.

We attribute the high novelty of Sc2Mol arises from the two-step approach of generating initial scaffold based on VAE and refining it to atom-level resolution. In addition, we carefully note that novelty is often anti-correlated with distributional fidelity, as we mentioned in Reviewer EQJM [Q1].

In the MOSES benchmark [18] (Official benchmarking implementation repository, Appendix C.1), novelty is computed as the percentage of novel molecules among the valid and unique ones. So the percentage of overall “novel and valid” molecules among the whole generated molecules can be computed by multiplying the validity, uniqueness and novelty in the table, resulting in 58.7%. (as uniqueness is 100%)

Respectfully, FragFM authors.

We are very grateful to the reviewer for the careful and detailed feedback. We appreciate the time invested, and comments have been instrumental in refining our manuscript.

We include a common response to all reviewers at the beginning of our response to Reviewer weUv, re-emphasizing our core contributions.

[W1] FragFM outperforms the baselines, but the gains are marginal

We thank the reviewer for pointing this out. We agree that some metrics on MOSES/GuacaMol appear close. Though, we carefully want to recap that on MOSES (Tab. 1) and GuacaMol (Tab. 4), FragFM attains FCD of 0.58 on MOSES and an overall FCD score of 85.8 on GuacaMol, substantially outperforming baselines (FCD 1.00-1.95; GuacaMol 0.9-73.8). To our knowledge, it is the first flow- and diffusion- based model to surpass autoregressive models on MOSES, while also achieving state-of-the-art Filters and SNN with nearly perfect validity. However, as reviewer mentioned, the performance gap is not significant, reflecting the saturation of these molecular generative benchmarks, where many methods already perform at near-optimal levels.

This saturation of performance is why we introduced the more challenging molecular generative benchmark, NPGen, which comprises diverse natural product molecules with high complexity. In this benchmark, the performance gap between FragFM and other models is far more significant.

For the evaluation of functional distributions (NPScore and NPClassifier pathway/superclass/class), FragFM yields remarkably lower divergence compared to the previous methods:

• JT-VAE: NPScore 14.54-, Pathway 5.38-, Superclass 8.70-, Class 7.18-fold better results
• DiGress: NPScore 5.23-, Pathway 1.17-, Superclass 2.27-, Class 2.89-fold better results.

In case of FCD, FragFM achieved 3.04- and 1.53-fold better results compared to the JT-VAE and DiGress, respectively. Aside from the benchmarks, we emphasize that FragFM-generated molecules avoid the chemically implausible motifs or avoid sampling overly simple molecules often observed in baseline models (App. F.3, Fig. 20-24).

To strengthen the evaluation, we include ZINC250K, another commonly used molecular generation benchmark. NSPDK evaluates the similarity of molecules by comparing the shortest paths in their graph structures. With 500 steps, FragFM achieved the state-of-the-art performance on all metrics. Especially, NSPDK and FCD show 5.00- and 2.31-fold better results.

[Q1] Insights into FragFM's lower novelty in benchmark

We appreciate the reviewer’s point: while FragFM attains high validity and strong distributional metrics (e.g. FCD/Filters/SNN), its novelty is lower than some baselines. In small-molecule benchmarks, novelty is often anti-correlated with distributional fidelity [81, 82]. FragFM also shows similar trend, with high validity and low FCD, with lower novelty.

Furthermore, we would like to claim that higher novelty does not necessarily imply chemically meaningful structures. As illustrated in the first paragraph of App. F.2, Fig. 20-24, baselines that report higher novelty sometimes produce chemically implausible motifs.

However, we found that it is possible to modulate the sampling algorithm to generate more unseen molecules, fully utilizing the flexibility of stochastic bag strategy. We add experiments that modulate sampling by introducing temperature factor to the Boltzmann probabilities. Specifically, two temperature factors are used:

1. $T_1$: reweighting fragment’s logit within the in-bag transition kernel (Eq. 5)
2. $T_2$: reweighting fragments when constructing fragment bag in each euler step (Eq. 34)

A higher temperature indicates that the sampling probabilities are more uniform. As shown in the following table, higher temperature factors reliably increases the novelty while allowing the trade-off with FCD and validity.

[81] Mahmood, Omar, et al. "Masked graph modeling for molecule generation." Nature communications 12.1 (2021): 3156.

[82] Geng, Zijie, et al. "De novo molecular generation via connection-aware motif mining." arXiv preprint arXiv:2302.01129 (2023).

[Q2] KL divergence calculation on NP-likeness scores

The Kullback-Leibler (KL) divergence for the continuous NP-likeness scores is calculated using a standard non-parametric method following previous benchmarks (MOSES, GuacaMol).

The procedure is as follows:

1. Score Calculation: First, we compute the NP-likeness score for every molecule in the reference test set and our generated set, resulting in two one-dimensional arrays of continuous values.
2. Distribution Estimation with KDE: Since the scores are continuous, we cannot use a simple histogram-based KL divergence. Instead, we model the underlying probability distribution of each set of scores using Gaussian Kernel Density Estimation (KDE). This technique creates a smooth, non-parametric probability density function (PDF) for both the reference distribution and the generated distribution.
3. Numerical Approximation: To compute the KL divergence integral numerically, we first define a common evaluation range that spans the minimum and maximum scores observed across both sets. We then evaluate both PDFs at 1,000 equidistant points along this range to obtain two discrete probability vectors. To ensure numerical stability, a small value (1e-10) is added to the probabilities.
4. KL Divergence Calculation: Finally, we compute the KL divergence between these two discrete probability vectors using the entropy function from scipy.stats.

[Q3] Dataset split on NPGen benchmark

We thank the reviewer for this important question about data splitting strategy of NPGen benchmark. We chose a random split rather than a similarity-based one because it is essential for NPGen to serve as a benchmark for distribution learning. In other words, the primary purpose of the NPGen benchmark is to evaluate how well a model can learn the complex, underlying distribution of natural products. A fundamental requirement for this is that the training and test sets are i.i.d. (independent and identically distributed). Our evaluation metrics, such as the KL divergence of pathway, superclass, and class distributions, measure how faithfully a model's generated output matches the training data's distribution.

Similarity-based split strategies (e.g., scaffold splits) can violate the i.i.d. assumption. This is particularly problematic when curating natural-product (NP) datasets and evaluating functionality-driven distributions, as molecules’ functionality often related with its structure. This creates training and test sets with fundamentally different distributions, making the benchmark unsuitable for evaluating distribution learning.

To quantify this, we analyzed the distribution of biosynthetic pathways after performing both a random split and a scaffold split on NPGen, which is predicted by NPClassifier [56]. The results clearly show that a scaffold split introduces a severe distributional bias:

As the table shows, the random split maintains near-perfect consistency between the train and test sets (e.g., < 4% deviation for all pathways). In contrast, the scaffold split creates massive shifts. For example, Fatty acids are over-represented in the training set and severely under-represented in the test set, leading to an 85.2% relative deviation.

Employing a scaffold-split dataset could lead to a misleading assessment of a distribution learning model's performance. A high KL divergence score would be an expected artifact of the biased split itself, rather than a true reflection of the model's generative capabilities. For this reason, we used a random split to establish NPGen as a fair and suitable benchmark for measuring a model's ability to learn the distribution of natural products.

Dear Reviewer EQJM,

We wanted to gently remind that the end of author-reviewer discussion period is approaching.

We appreciate the reviewer’s feedback and their recognition of FragFM as a novel framework for molecular graph generation as well as highlighting NPGen as addressing an unmet need and opening new avenues to evaluate the bio-compatibility of molecular generative models.

We have conducted additional experiments and revised the manuscript to address your concerns:

• [W1] we re-emphasized the original results to clearly show a significant performance improvement over existing atom-based baselines in NPGen, with our performance being at least twice as good as the previous state-of-the-art (SOTA) model for most distribution-learning metrics. Furthermore, we note that for the standard molecular generative benchmarks, which is saturated, FragFM clearly approached its performance toward the optimal limit of each metric in FCD, Filters, and SNN.
• [Q1] we provided a deeper explanation of novelty and included temperature-adjusted sampling experiments to steer generation toward novel molecules.
• [Q2]–[Q3] we expanded explanations and analyses for the NPGen benchmark.

Your feedback was very helpful: it prompted us to (1) emphasize FragFM’s larger margins over baselines, (2) add a steering strategy for sampling more novel molecules, and (3) improve the readers' understanding of the NPGen benchmarks. We will integrate these clarifications into the manuscript text.

We also emphasize that we organized our core contributions in Reviewer weUv [Common response to all reviewers].

We hope the added experiments and updates address your original concerns. If you feel these revisions strengthen the paper, we’d be grateful if you could reflect that in your rating; and if anything remains unclear, we’re happy to follow up.

Respectfully, FragFM authors.

We sincerely appreciate your confirmation that our rebuttal has addressed your concerns. If you believe that the discussion has led to meaningful improvements, we would be grateful if you could take this into account when you make your final evaluation. Thank you again for your valuable feedback and time.

We sincerely thank the reviewer for a thorough, thoughtful review regarding the clearity of the manuscript and experimental designs.

We include a common response to all reviewers at the beginning of our response to Reviewer weUv, re-emphasizing our core contributions.

[W1] Manuscript's clarity is impeded by a structural flaw

For the implementation details about hyperparameters and coarse-to-fine autoencoder, please refer to our reply for Reviewer, A3wR’s [W1,Q1]. As many reviewers pointed out limited explanation for coarse-to-fine autoencoder in the main text, we will revise the manuscript to include more information about the coarse-to-fine autoencoder.

Furthermore, we agree that our manuscript should place greater emphasis on the stochastic fragment bag strategy, which is a core contribution to our model's efficiency and performance. To address this issue, we will revise the relevant subsections or paragraphs in Section 3 to give this strategy the prominence it deserves. In particular, as sections 3.3 and 3.4 mainly introduce the fragment bag strategy for training and generation, we will emphasize the stochastic fragment bag strategy as follows(we have highlighted the revisions with in bold):

• In line 102-105: Starting from fragment graph notation in section 3.1, we elaborate on ... fragment-level graph based on stochastic fragment bag strategy in sections 3.3 and 3.4.
• In line 149: Stochastic Fragment Bag Strategy: Parameterization and Info-NCE Loss
• In line 167: 3.4 Generation Process with Stochastic Fragment Bag Strategy

[W2,Q2] A significant gap exists between the paper's proposed contributions and its experimental validation

Reviewer’s comments on weaknesses and questions

The manuscript fails to provide … of the central claims of the work.

FragFM introduces the fragment bag strategy … rare fragments?

Thank you for this thoughtful comment. Our main goal in FragFM is to show that moving from atom-wise diffusion to a fragment-level flow framework brings three practical advantages:

1. Scalability to complex molecules
2. Much faster sampling while preserving validity
3. Stronger and more flexible property control

The experiments included in the paper were chosen to demonstrate these three points side-by-side against atom-wise baselines.

Though, we fully agree that readers will also want direct evidence that the fragment-bag strategy copes with challenges arising in a fragment-based framework, such as the long-tailed distribution of fragment types and the need to generalize to novel fragments that never appear during training. We have therefore run an additional set of analyses and will incorporate these results into the revision.

On rare/long-tail fragments. Handling the long tail (seldom showing fragments) is essential, especially when the number of fragments are large (e.g. NPGen has 133K fragments). In theory, our stochastic bag is an unbiased estimator of the full-fragment vocabulary transition kernel (Eq. 7, our reply on Reviewer weUv’s [W1,Q1]), so rare fragments are not excluded by design. More specifically, to verify that rare fragments appear in the generated molecules at the same rates as in the training data, we conducted a rare-fragment recovery analysis. By defining rare as fragments that appear $k\\in \\{ 1,\\ldots,5 \\}$ in the data, we compare their occurance shares between the training and 100k generated molecules (reported per-$k$ bin). (i.e., total occurrences of k-frequency types divided by total fragment occurrences). The measured ratios closely resembles the training distribution, indicating that our bag strategy can effectively generates the molecules with rare fragments.

Novel fragments. Generation with test set fragment bags is reported in the Tables 6 and 7, directly evaluating generalization to fragments unseen during training. For clarity, we also enumerate the counts: in MOSES, the test split contains 2,588 unseen and 7,210 seen fragment types (26% unseen); in GuacaMol, 22,962 unseen and 40,646 seen fragment types (36% unseen). Under these settings, FragFM maintains state-of-the-art validity and distributional metrics, demonstrating robust generalization to novel (unseen) fragments.

[Q1] Validation of autoencoder for its lossless reconstruction ability

Reviewer’s comments on questions

The authors claim that this work achieves lossless reconstruction … FragFM could achieve lossless reconstruction?

On Line 124, the authors claim that the single continuous latent vector $z$ encodes the committed connectivity details ... validity.

We sincerely thank the reviewer for the careful reading and the question. For clarity, we distinguish reconstruction accuracy from validity. Accuracy means that, given a fragment graph and its latent $z$, the decoder recovers the exact same molecule (e.g., from fragments A*, *BC*, D* (for ABCD, where * indicates the possible junction)). Validity only requires chemical feasibility—so DBCA would also count.

Our coarse-to-fine autoencoder uses Blossom matching algorithm (App. 3.2) to reconstruct atom-atom connectivity from fragment-level graph and $z$. The gap to 100% generation validity (relative to JT-VAE) mostly arises from invalid combination of fragments (for example, A* , *BC*), which cannot be reconstructed to any atom-level graph. We further conducted analysis about the role of $z$ and validities of JT-VAE and FragFM.

On the role of the latent $z$ in autoencoder accuracy. To show that latent $z$ encodes atom-level connectivity information well, we conduct an ablation study on $z$ for the reconstruction accuracy. Our autoencoder achieves almost perfect reconstruction of all test sets, whereas decoding with random $z \\sim \\mathcal{N}(0,I)$ does not. As the size of the molecules in the NPGen benchmark larger, the discrepancy in accuracy rises, empirically showing that the latent $z$ encodes almost-full information of the atom level connectivity information.

Since reconstruction accuracies do not achieved 100.0% in all dataset, we will tone-down the phrase “lossless reconstruction” into “nearly lossless reconstruction”.

On JT-VAE’s 100% validity compared to FragFM. JT-VAE enforces validity by design: during decoding it masks fragment choices that would violate chemical feasibility and prunes infeasible assemblies, restricting the search to valid molecules. In contrast, diffusion/flow-based models update many nodes and edges simultaneously, making hard-encoding of rules difficult; consequently, achieving literal 100% validity is inherently harder (e.g., DiGress ~85–86%, DeFoG ~92–93%; Tab. 1). Our fragment-level approach reduces this complexity and attains >99% validity on MOSES in a purely data-driven manner (no rule-based masks) while delivering stronger distributional metrics, e.g., Filters 99.1, FCD 0.58, SNN 0.56 than JT-VAE (Filters 95.0, FCD 1.00, SNN 0.54).

[Q3] Choice of Info-NCE loss

We adopt Info-NCE because it is the only objective we have found that can model $p_{1|t}(x|X_t)$ without bias. In our method, the density ratio

f_\\theta(x_{t}; X_t,t)=\\frac{p_{1|t}(x_{t}|X_t)}{p_{1}(x_t)},

is modeled directly (Eq. 5) and is independent of the bag size $N$ [44]. Consequently, although the model is trained with a fixed fragment-bag size $N$, it can be applied with variable bag size during inference.

Computing a full soft-max over the entire fragment library $|F|\\approx 10^{4\\sim5}$ is infeasible, whereas Info-NCE restricts computation to the $N$ fragments in the fragment bag while remaining an unbiased estimator of the full objective for moderate $N$.

[Q4] Node, Edge diffusion

Factoring node- and edge-state updates is a standard practice in discrete diffusion/flow models[3,28,32] and has proven empirically sufficient: despite the independence assumption, our method still outperforms atom-wise baselines on MOSES and NPGen (Tab. 1,2).

Coupling every node flip with all incident edge flips would force a non-linear trajectory in Eq. 2, leading to intractable Euler updates and invalidating the closed-form rates in Eq. 3. Ensuring graph consistency under such coupling would also require additional constraints that are infeasible. While joint node-edge transitions are an interesting direction for the future, to the best of our knowledge no tractable implementation has been demonstrated so far. We therefore retain the independent formulation, which is both computationally feasible and empirically validated.

Our fragment-based approach is inherently better suited to mitigate this dependency issue than atom-level models. In atom-level diffusion, independent node and edge updates frequently violate local valency rules, leading to chemically invalid states, as shown in benchmarks (Tab. 1,2,7). FragFM mitigates this problem by design: since each node is already a chemically coherent fragment, the most critical local dependencies are implicitly preserved within the node itself. This structural constraint is why our model achieves over 99% validity in a purely data-driven manner, directly demonstrating that the information loss from the independence assumption is minimal.

Dear reviewer NJHQ,

We wanted to gently remind you that our discussion period is end soon.

We appreciate the reviewer’s recognition of our experiments across multiple benchmarks, including the conditioning and sampling-efficiency studies. We also thank the reviewer for acknowledging NPGen as a useful benchmark that contributes to both the machine-learning and drug-discovery communities.

Thank you again for the detailed, specific review. We added explanations and analyses that directly address your concerns:

• [W1] clearer explanation of the method
• [W2,Q2] rare/novel fragment coverage (recovery analysis and test-bag generalization)
• [Q1] coarse-to-fine autoencoder reconstruction accuracy
• [Q3] Explanation about the choice of Info-NCE loss
• [Q4] Reasoning about independent node and edge diffusion

Adding on our initial response to [Q3], we also had included additional theory and experiments (see Reviewer weUv [W1,Q1] and Reviewer A3wR [Q2]), with a detailed study of fragment-bag size. We will add these clarifications into the main text for readability.

We also emphasize that we organized our core contributions in Reviewer weUv [Common response to all reviewers].

We acknowledge that the initial submission could have better highlighted each component individually. Since we have now added this material, we hope the revisions sufficiently address your concerns and kindly ask you to reconsider revising the evaluation of the paper.

Respectfully, FragFM authors.