Compositional Flows for 3D Molecule and Synthesis Pathway Co-design

We appreciate the reviewer for volunteering their valuable time and providing insightful feedback to our paper. We are addressing their questions one by one in our response below.

W1. Evaluation of the diversity and chemical properties.

To further address the reviewer’s suggestion, we now compare sampling efficiency, diversity and other chemical properties of our method (CGFlow) against the 2D baseline (RxnFlow) on the first 5 LIT-PCBA targets. We define high-scoring modes as those with QED > 0.5, Vina < -10 kcal/mol for all pockets except FEN1, which we use -7 kcal/mol, and mode similarity < 0.5.

The table below shows the number of unique high-scoring modes identified after sampling 10k molecules:

CGFlow consistently outperforms RxnFlow in sampling efficiency, discovering 4.7x more diverse modes. Since mode diversity increases experimental success, this improvement highlights the practical advantage of our approach.

We further report the full sampling trend for ADRB2 and the average properties of the top 100 diverse modes. Diversity here is computed without similarity-based filtering to avoid artificial inflation. Our results confirm CGFlow's superior efficiency in discovering diverse modes with good QED and Vina scores.

We thank the reviewer for their comment, which motivated us to further highlight the strengths of our approach.

W2. Concern about computational cost.

We kindly refer the reviewer to our response to Reviewer Z8PT (Weakness 2) for further clarification on computational cost.

W3. Concern about limited search space due to the brick-and-linker formulation.

We appreciate the reviewer highlighting this important limitation regarding our synthesis approach. Indeed, as the reviewer pointed out, our current brick-and-linker formulation restricts the explored chemical space to linear molecules by excluding nonlinear reactions, such as ring formation. However, this constraint can be substantially mitigated by expanding the chemical search space through incorporating a larger building block library. For example, V-SYNTHES [1], a pioneering work in exploring Enamine REAL Space using a brick-and-linker strategy, successfully achieved a notable hit rate of 33%.

Q1. Question about the size of state space and unique building blocks explored during training.

We estimate the sample space according to the number of synthetic steps: $10^{11}$ molecules with a single reaction step, $10^{17}$ molecules with two reaction steps, and $10^{23}$ molecules with three reaction steps. In our experiments, we employed up to two reaction steps according to Enamine REAL, and the state space size is similar to RGFN (up to 4 steps with 8,350 blocks) and SynFlowNet (up to 3 steps with 200k blocks).

Additionally, we analyzed the number of unique building blocks (BBs) explored during training across the first 5 LIT-PCBA targets. Our model explored an average of ~55,000 unique BBs within 1,000 training iterations with a batch size of 64. This demonstrates a significantly broader exploration compared to SynFlowNet, which reported exploring ~15,000 unique BBs during 8,000 training iterations with a batch size of 8.

Q2. clarification as to the training details of the state flow model.

Yes, you are absolutely correct that during training, the state flow model encounters "partial" CrossDocked molecules. Specifically, we decompose each molecule into up to three fragments using 38 bimolecular Enamine synthesis protocols defined by reaction SMARTS. We then randomly sample a fragment ordering, matching the fragment introduction schedule used in the Compositional Flow model (e.g., fragment A at $t=0$, B at $t=0.3$, etc.). A random time step is sampled so that at earlier steps, the model sees "partial" structures. This design enables the state flow model to learn realistic fragment docking conformations aligned with the fragment selection process of compositional flow.

Importantly, the decomposition is not intended to recover purchasable Enamine synthons but to expose the model to chemically meaningful substructures for learning protein-pocket conformations. While guided by Enamine protocols, exact synthon matching is unnecessary, allowing us to use any molecules for pose prediction training. We will clarify these points in the revised appendix.

Reference:

1. Sadybekov, Arman A., et al. "Synthon-based ligand discovery in virtual libraries of over 11 billion compounds." Nature 601.7893 (2022): 452-459.

We highly appreciate this reviewer’s constructive feedback and insightful suggestions. We would like to clarify and address all of these points to the best of our ability in the response below.

W1: Lack of ablation studies about time scheduling of state flow model

Thank you for the valuable suggestion. We conducted an ablation study to assess the effect of time scheduling in state flow training, comparing three settings: partial (partial overlap of synthon denoising), no overlap (strictly autoregressive), and till end (all synthons denoised until $t=1$).

We compared the average local-optimized Vina docking scores across different training iterations for the ALDH1 target below:

CGFlow’s overlapping noise scheduling, where positions are refined as synthons are added, clearly outperforms conventional autoregressive approaches (no overlap).

W2 & W3: Lack of comprehensive evaluation (e.g., CrossDocked2020) and SBDD baselines.

Following your suggestions, we evaluated CGFlow on CrossDocked2020 against established SBDD baselines. Using the same conditional objective and proxy setup as TacoGFN and RxnFlow, we generated 100 molecules per pocket in a zero-shot manner without an additional optimizing process for test targets. We varied the reward exponentiation parameter β (Low: U(1,64), Medium: U(32,64), High: U(48,64)) to balance exploitation and exploration for sampling.

CGFlow reduces Vina from -8.85 (RxnFlow) to -9.38 (CGFlow-high beta), outperforming all baselines. It also yields the highest QED scores (0.72–0.74) and highest AiZynthFinder success rate (62.2%) compared to all baselines, underscoring the practical benefits of synthesis-aware generation. CGFlow shows consistent synthesis success rate across both CrossDock (55.0%–62.2%) and LIT-PCBA (53.1%) benchmarks.

W4. Concerns about reward hacking by generating larger molecules.

To address this concern, we conducted additional experiments on the first five targets, restricting heavy atom count (HAC) to 40. CGFlow still outperforms RxnFlow in Vina score (-10.94 vs -10.46) with comparable HAC (29.63 vs 29.37). Moreover, CGFlow achieves the highest ligand efficiency (0.375) - computed by Vina / HAC, confirming that our binding affinity gains stem from the 3D co-design strategy rather than molecule size.

W2 / W4: Measurement of training efficiency / The vina improvement over RxnFlow is marginal

We kindly refer the reviewer to our response to Reviewer ScwC (Weakness 1) for experimental results on training efficiency - where we show CGFlow discovers 4.7× more diverse modes than RxnFlow. We note that optimization of docking scores is restricted by the saturation of the pocket's binding interactions. At that point, discovering more diverse binding modes becomes more important to maximize the success rate of practical applications.

W4. Degeneration in Success Rate (synthesizability) and synthesize steps

The small drop in AiZynthFinder synthesizability arises from our transition from reaction-based generation (RxnFlow) to a synthon-based (brick-and-linker) approach. Reaction-based generation often halts prematurely if a state molecule lacks any reactive functional groups, while the synthon-based method can easily construct molecules with longer synthetic trajectories. We emphasize that our approach use the building block and synthesis reactions from Enamine REAL and xREAL, known for a wet-lab synthetic success rate of 80%.

W2/Q1. Lack of evaluation in geometrical properties and questions regarding generated poses.

We evaluated various geometrical properties and Vina score of the generated poses of the top 100 molecules of local-optimized Vina optimization across 3 seeds.

We sincerely thank the reviewer for their thoughtful and detailed evaluation, and for recognizing the significance of our framework contribution and its application to 3D molecule and synthesis pathway co-design.

W1. More analysis of cases where the method performs poorly or limitations in certain chemical spaces.

Our synthesis-based action space based on brick & linker synthons does not yet explore certain synthetic pathways such as ring-forming reactions, and nonlinear synthetic pathways, as pointed out by Reviewer ScwC. Our method predicts the 3D binding poses of intermediate states. This synthon-based generation is chosen to prevent the atoms from intermediate states being lost from reactions or forming new rings - which would degrade the accuracy of pose prediction for intermediate states.

Moreover, our pose prediction module is trained on the CrossDocked2020 dataset, in which binding pockets are extracted with a distance cutoff based on the reference ligand structure, resulting in an inherently biased pocket to its reference ligand. This can negatively impact pose prediction accuracy if generated molecules optimally bind to a different subpocket from the reference ligand or are larger than the reference ligand. This issue arises during reinforcement learning-driven exploration of diverse chemical spaces. Addressing this limitation would require employing an unbiased pocket structure (e.g., using a center-based cutoff) or a full-protein structure, and we highlight this as a consideration for real-world application of our method.

Finally, our training of poses is currently biased towards Vina docking poses from the CrossDocked2020 dataset. A promising direction to remove this bias is training the model on an experimental structure dataset.

W2. A more explicit discussion of the implementation challenges and computational overhead.

Thank you for this suggestion! We will incorporate the following details in our final manuscript. The state-flow model (i.e., the pocket-conditional pose predictor) was trained for 100 epochs using a batch size of 32 on 4 L40 GPUs (48GB), taking a total of 18.4 hours. Importantly, as the state-flow model is trained on the CrossDocked dataset and can be reused across different test pockets, it incurs only a one-time computational cost. We also plan to release the model weights, so users will only need to train the composition-flow model tailored to their custom reward function and target.

The composition-flow model is trained individually for each pocket for 1,000 steps with batch size of 64 and 80 flow matching steps. The training with GPU-accelerated docking takes 12-20 hours (depending on targets) on a single A4000 GPU (16GB). We find this computational requirement accessible for most practical drug discovery campaigns. Furthermore, the composition-flow model can also be trained in a pocket-conditioned manner (Please see SBDD benchmark in our response to Reviewer eDTU) to sample molecules for any target pocket in a zero-shot manner, making model training a one-time cost in this setting.

The sampling time of our model is 0.05~0.20 seconds/molecule depending on the choice of number of flow matching steps. We further analyzed how the number of flow matching steps impacts performance using the ALDH1 target with the Vina reward. Performance slightly improves with increased flow matching steps and saturates around 40-60 steps. We attribute this marginal improvement to the fact that the pose prediction module’s primary role is providing a spatial context between intermediate molecules and the pocket; thus, extremely precise pose predictions have limited additional impact on model decisions.

Theoretical Claim & Question

We thank the reviewer for this insightful observation. Indeed, our current theoretical analysis leverages determinism through fixing the random seed primarily for analytical clarity. However, this leads to deterministic noise sampling for the initial synthon states; for example, when C(=O)[*] and [*]NC are sequentially added, their initial coordinates are identical. To maintain determinism while preventing synthons from receiving the same initial states, we can use the molecule size or a hash of the molecule to set distinct random seeds across the generation process.

We acknowledge this limitation explicitly and highlight extending our theoretical framework to handle fully stochastic environments, using Expected Flow Networks [1] as a promising direction for future research.

Reference:

1. Jiralerspong, Marco, et al. "Expected flow networks in stochastic environments and two-player zero-sum games." ICLR 2024.