We thank the reviewer for finding our work adaptable leading to promising results especially for 3D structures. We also appreciate that our new benchmarks were found to be more aligned with practical applications of molecule design and drug discovery.

We thank the reviewer for their time and understand that novelty of a method can be highly debatable. For this reason we will highlight aspects of our work that we find novel below. If there are more specific questions we are happy to answer them

Architecture Novelty

The standard DiT block takes in a singular input tensor, H. To enable the modeling of 3D molecules composed of several continuous and discrete data modalities, we made several critical changes to the architecture, which we refer to as fused-DiT (f-DiT).

• As described in Sec. B.1.2, our fused DiT blocks take in (X, H, E, C) for the 3D structure and discrete atom, bond, and charge types. These features are first fused and aggregated in a message-passing-like operation. The multi-head self-attention is then applied to these fused features. We note that if the attention is applied to the respective inputs as done in the standard DiT, the model does not converge and outputs 100% invalid molecules with unrealistic structures
• For the feed-forward part of the f-Dit block, the processed fused features are aggregated along all pairs of nodes to create new bond features, and they are hit with a linear layer to create new atom and charge features. From here, we apply independent feedforward and adaptive layernorm operations for each data modality (excluding structure since the structure is only used as an input to the f-DiT block) to create updated features for all discrete data modalities.
  • In short H’, E’, C’ = f-DIT(X,H,E,C)
  • We emphasize that the only operation maintaining equivariance and updating the structure is the series of EGNN single layers. In comparison, if we replace the f-DiT operation with standard EGNN non equivariant feature updates, we see a drop in Validity of almost 70%, as shown in Table 1. From our understanding, this is the first work to obtain such strong results with a simple EGNN-based architecture since EDM + Openbabel.

Technical Novelty

Our novelty is further rooted in our applications, evaluations, and analysis, which include the introduction of new benchmarks and the reintroduction of a conditional structure generation task.**

• Figure 3 introduces molecule size as a new component to unconditional benchmarking in which the Megalodon significantly outperforms EQGAT-diff. We demonstrate this performance can be improved with further scaling of our architecture. The ability to generate 49x more valid and stable molecules compared to prior SOTA is a significant result, given that both are trained with identical data and diffusion parameterizations.
• Table 2 demonstrates that off-the-shelf 3D molecule generation models cannot be used for conformer generation. In contrast, Megalodon can and surpasses strong conformer baselines due to the use of a compounded time-dependent noise scheduler described in line 297. This is a novel and quite surprising finding, as we expected all unconditional molecule generation models to be able to conditionally generate structure as they are trained with independent structure and discrete denoising.
• Furthermore, when compared to GeoDiff, which also uses an EGNN-based architecture with identical diffusion parameterization, we demonstrate that unconditional generative pretraining is extremely beneficial in generating better structures in 10x fewer sampling steps. In other words, learning how to generate the 2D discrete components improves the ability to generate accurate conformers.
• We demonstrate that with the diffusion objective, Megalodon is capable of generating conformers—that is, molecules very close to their local minima of the ground truth energy function. The median relaxation energy drop of 3.17 kcal/mol approaches the threshold of 2.5 kcal/mol, which is often considered the thermodynamically relevant interval. Furthermore this is 2-10x better than prior methods. This proximity emphasizes the potential practical value of our method. In contrast, we showed that with the Flow Matching objective, the energy drop is an order of magnitude larger, highlighting a significant and valuable difference for readers.
• Subsequently, we uncover an efficiency and accuracy tradeoff between FM and DM for 3DMG. FM yields more valid models and can be constrained to very few sampling steps, whereas DM exhibits an order of magnitude better structure measured by the molecular energy.

Overall, we are the first to perform a comprehensive analysis of the interplay between the 2D graph and 3D structure during molecular generation in previous methods and generative frameworks(Diffusion vs Flow Matching) and offer potential solutions to improve the dependency between the modalities.

We thank the reviewer for highlighting key contributions of Megalodon, including its ability to advance 3D molecule generation (3DMG) into practical applications with improved performance on larger molecules and practical structural benchmarks. We’re glad that our in-depth analysis of the 2D-3D interplay in molecular generation was found to be thorough and valuable.

We are happy to provide more details on the provided questions below.

Q1. How is equivariance preserved?

Thank you for pointing out the inconsistency in Figure 1, and we apologize for any confusion this may have caused. There is a skip connection seen in Eqn. 17, which we will add to the camera-ready figure. In our architecture, the DiT (Diffusion Transformer) block operates on invariant features such as atom types, charges, bond types, pairwise atom distances, and coordinate norms. The output of the DiT block consists of an updated atom (H), charge (C), and bond (E). In practice, C is concatenated to H but we write it explicitly here to be more clear.

• H’, E’, C’ = f-DIT(X,H,E,C)
• X’ = EGNN_X_Only(X, H’, E’, C’) #only x update see Eqn. 17

The structure block then uses these updated atom and bond features to update the coordinates. Since we subtract the center of mass (CoM) from the coordinates, all features processed by the DiT block are invariant under rotations and translations. Therefore, the DiT block updates only invariant features.

Our structure block follows the Equivariant Graph Neural Network (EGNN) architecture [Satorras et al., 2021 https://arxiv.org/pdf/2102.09844], as illustrated in the formula in Figure 1. This design ensures that the overall architecture is equivariant, meaning that any rotations or translations applied to the input coordinates result in corresponding rotations or translations in the output coordinates. The skip connections from the input 3D coordinates to the structure block are implicit in the EGNN framework.

Q2. How is the 2D graph provided to the diffusion model?

To provide the 2D graph to the diffusion model, we supply the atom types (H), bond types (E), and charge types (C) as fixed inputs, while the coordinates (X) are generated by the model.

In standard diffusion and flow matching models that directly operate on bonds, the initial inputs H,E,C,X are generated from a prior distribution. At each diffusion step t, the model takes the noised versions of these variables Ht,Et,Ct,Xt​ and predicts the ground truth values.

To enable conditional generation, we modified the training process so that for a fraction of the time the model is conditioned on the ground truth H,E,C, which are supplied to the model as one-hot vectors, and we use RDKit to compute adjacency matrix, bond orders, atom types and formal charges. This means that during conditional generation, the model receives the fixed 2D graph (represented by H,E,C) and the noised coordinates Xt​, and it predicts the denoised coordinates X.

In summary, the model takes the fixed atom types, bond types, and charges as inputs and generates the corresponding 3D coordinates, effectively performing conditional generation based on the provided 2D molecular graph.

If you have any future concerns or questions, we will be happy to address them.

Thank you very much for your effort in reviewing our paper and your feedback. We appreciate the depth of the review and give our best effort to answer all questions.

W1: Missing training and evaluation details missing

We provide links to the corresponding code bases from which we used identical training setups for the diffusion and flow matching models. We do this to create a true 1:1 comparison with both EQGAT-Diff and SemlaFlow, with the only change being the architecture we describe in Appendix Section B, which includes all key equations and hyperparameters. This allows us to use the same data loaders and generative model setup as each method, which is important as SemlaFlow filters out all molecules with greater than 75 atoms. We also provide all evaluation details and use standard DDPM SDE and flow matching ODE sampling as done in prior work.

Outside of the training and evaluation details that we clarified above, are there any sections of the paper that were not clear?

W2-A: Architecture Novelty

Our work is more than applying diffusion (DM) and flow matching (FM) for unconditional molecule generation.

First, we want to clarify that DiT Block, while used for images in the original paper, does not have any specific inductive bias for images. As discussed here https://www.wpeebles.com/DiT, the DiT block is just a standard transformer block with an adapted layer norm to handle the conditioning input. It was used for latent image diffusion, but that use case has no bearing on how we augment it for molecule generation.

We emphasize that we deconstructed molecule generation into simultaneous equivariant structure prediction and non-equivariant discrete data prediction (atom, bond, and charge types). This is reflected in our architecture design, which uses an augmented DiT to model the discrete data and then simple EGNN layers to update the structure. Figure 1 illustrates how the model comprises N segments of a DiT block followed by a simple structure update. We chose this architecture to enable better modeling of discrete data via the multi-head self-attention as the discrete data accuracy determines molecule stability and validity.

The standard DiT block takes in a singular input tensor, H. To enable the modeling of 3D molecules composed of several continuous and discrete data modalities, we made several critical changes to the architecture, which we refer to as fused-DiT (f-DiT).

• As described in Sec. B.1.2, our fused DiT blocks take in (X, H, E, C) for the 3D structure and discrete atom, bond, and charge types. These features are first fused and aggregated in a message-passing-like operation. The multi-head self-attention is then applied to these fused features. We note that if the attention is applied to the respective inputs as done in the standard DiT, the model does not converge and outputs 100% invalid molecules with unrealistic structures.
• For the feed-forward part of the f-Dit block, the processed fused features are aggregated along all pairs of nodes to create new bond features, and they are hit with a linear layer to create new atom and charge features. From here, we apply independent feedforward and adaptive layernorm operations for each data modality (excluding structure since the structure is only used as an input to the f-DiT block) to create updated features for all discrete data modalities.
  • In short H’, E’, C’ = f-DIT(X,H,E,C)
  • We emphasize that the only operation maintaining equivariance and updating the structure is the series of EGNN single layers. In comparison, if we replace the f-DiT operation with standard EGNN non equivariant feature updates, we see a drop in Validity of almost 70%, as shown in Table 1. From our understanding, this is the first work to obtain such strong results with a simple EGNN-based architecture since EDM + Openbabel.

We thank you for your feedback. While we disagree that our work lacks novelty due to enabling multi variable message passing into the DiT block which was not a minor modification, high quality unconditional generation and high quality conformer generation in and same model for the first time, and also introducing new interpretable benchmarks, we thank you for asking us questions and making our work stronger.

As for diffusion and flow matching we know they have been used before but we are the first to study them together using the same architecture and biological task at the same time. We find our findings of a physical accuracy vs efficiency trade off between diffusion and fm models quite interesting and something potentially meaningful for other applications of these frameworks outside of small molecule design.

Overall we have answered all questions and addressed all concerns that can be addressed in the rebuttal period. We understand novelty is highly debated and value your feedback.

I appreciate the authors' rebuttal and the clarifications. I also acknowledge the good empirical results on GEOM-drugs. Although this does not necessarily come as much of a surprise given it has more parameters than the compared baselines. I still find the paper lacks novel technical contributions: the architecture (modulo minor modifications), the data representation, and the generative models utilized have already been used in many similar applications. For these reasons, I keep my current rating.