Disentangled autoencoding equivariant diffusion model for controlled generation of 3D molecules

Domain: biology
Published: 2026-04-03
Editorial outcome: major_revision

Abstract
Controlled generation of 3D molecules is highly valuable in drug design, enabling targeted modifications that preserve core structure and binding-relevant geometry while improving developability-related properties. While equivariant diffusion models achieve state-of-the-art de novo 3D molecule generation, reliably controlling multiple molecular properties remains challenging. A key limitation is structural: A diffusion models typically lack an explicit latent space for targeted manipulation. We propose a semantics-guided equivariant autoencoding diffusion model that learns a disentangled semantic embedding of 3D molecules via an auxiliary encoder, to achieve fine-grained control over the generative denoising process. This semantic embedding enables efficient retrieval, random generation and controlled generation. By directly manipulating the embedding, we effectively steer the molecular generation toward desired compositions, shapes and physicochemical properties, and further enhance the generation quality with retrieval-augmented generation (RAG) using the embedding as the query. Importantly, the disentangled embedding offers significant advantages for joint manipulation of multiple properties. Experiments demonstrate precise and data-efficient property control while preserving non-targeted properties.

Introduction
Computational design of molecules has many applications in drug design [1, 2] and protein engineering [3–5]. The 3D geometry of molecules holds significant implications for their properties and functions, such as quantum chemical properties [6, 7], molecular dynamics [8–10], and interactions with protein receptors [11, 12]. Recently, generative modeling of 3D molecules has been an active area of research. To properly represent the 3D geometry, some methods learn a joint distribution over the geometric features [4, 13]. Others directly operate on the 3D molecular space [14]. Specifically, equivariant diffusion models have achieved state-of-the-art results in generating high-quality 3D molecules [15–17], and also in design tasks with external conditioning, such as ligands conditioned on protein receptors [18, 19] and linkers conditioned on molecular fragments [20]. Diffusion models belong to the class of generative models designed to approximate the data distribution by iterative denoising from corrupted data [21, 22]. Unlike previous methods for molecule generation that rely on autoregressive generation, diffusion models perform simultaneous refinement of all elements, such as atoms, charges and bonds, in each denoising iteration. This collective refinement approach enables them to effectively capture the underlying data structures, making them well-suited for structured generation such as 3D molecules. Several recent studies further incorporate auxiliary information, such as bond connectivity, to improve the generation quality [23–25].

However, the majority of these works focus on de novo generation, i.e., generating random 3D molecules with no or limited control over their compositions, structures and properties, which we collectively define in this work as “semantics”. In real-world applications such as drug design, it would be desirable to ensure that the generated molecules contain certain attributes through explicit control of the generation process. Some previous works focus on the controlled generation of 2D graphs [1, 26–32], while controlling 3D molecule generation remains a challenge due to the complexity of the 3D molecular space and the need to preserve the 3D geometry. Some recent studies achieve conditioning on selected molecular properties [15, 33, 34] or 3D shapes [35]. These methods, despite being effective on the designated tasks, only control a narrow portion of the generation process, and cannot be easily generalized to multiple or novel conditions without training a new conditional generative model. On the other hand, it is uncertain how conditioning affects the interactions among properties due to their complex interplay, and manually disentangling them is non-trivial and time-consuming. For example, a lead compound may require higher water solubility and easier synthesis. Meanwhile, the backbone composition, the 3D shape and subsequently the protein binding affinity of the lead compound should also be preserved. As these objectives do not necessarily align, the generation process that fulfills one objective may potentially interfere with another, and naively conditioning on multiple objectives without proper regularization may fail to find a direction that satisfies all. To overcome these limitations, it is thus essential to implement a refined control mechanism that both captures the comprehensive context of molecular compositions and shapes, and automatically disentangles the intricate interdependencies between the properties.

Controlled generation with diffusion models in particular poses another challenge as these models lack an explicit latent space to operate on. The latent diffusion model [36] is trained on the low-dimensional latent space of another generative model, but the diffusion model’s own latent space is still implicit and difficult to control. Thus, controlled generation is usually achieved by external guidance on the denoising process. For example, classifier guidance [37] and classifier-free guidance [38] can generate random samples conditioned on certain properties. These methods have high demands on the availability of labeled data, and generalize poorly to new or multiple objectives. Some studies utilize external, pre-trained molecular representations as conditions [39]. Though such representations could be highly informative, there is no theoretical guarantee of the effectiveness of their control, and they cannot be easily manipulated. Thus, for data-efficient control of multiple molecular properties, we could leverage an unsupervised latent space that captures the complete semantic information about molecular structures and compositions and use it to guide the generation process of diffusion models.

In this work, we propose MolDiffdAE, a disentangled autoencoding equivariant diffusion model controlled by an unsupervised, higher-level semantic embedding of the molecule, inspired by the idea of semantic guidance of diffusion models for controlled generation of images [40–42]. The semantic embedding is co-trained with the diffusion model and is set up to have three important features that address the aforementioned challenges: (1) maximal mutual information with the input that guarantees complete semantic information about the molecule; (2) strong and precise control over the generation process of the diffusion model that effectively translates the semantics into the 3D molecular space; (3) disentangled semantic space through distributional regularization, which empirically enables conditioning on multiple properties with minimal internal interference, yielding Pareto-efficient control directions. By directly operating on the semantic embedding, we could control and manipulate the compositional and structural features of the generated molecules, as well as higher-level properties. We demonstrate the effectiveness of this method in two types of generative tasks: (1) semantics-guided generation, where we use predefined semantics to generate random molecules that contain the desired semantic information; (2) property manipulation, where certain properties of known source molecules are modified through altering the respective dimensions of the semantic embedding. We then evaluate whether the generated molecules exhibit the specified property(ies). We summarize our main contributions as follows:

• We perform a comprehensive analysis of controlled 3D molecule generation and manipulation using diffusion models. We also establish the benchmark tasks and evaluation metrics for the subject. We prove that the semantic embedding learned by our model has a theoretical guarantee of maximal mutual information with the molecule, and demonstrate effective control over more than 200 molecular properties.

• We transform the challenging task of controlled molecule generation using 3D equivariant diffusion models, often requiring sufficient labeled data or accurate classifiers, to the more data-efficient task of latent space optimization (LSO). Thus, our model could be flexibly applied to challenging molecular design tasks where labeled data is scarce and complex relations are present among the objectives. We show that the semantic embedding also allows fast querying of the molecular space and retrieval-augmented generation (RAG) to further enhance the generation quality.

• Compared to existing baseline methods, we report the highest performance on multi-objective property manipulation, a task hindered by the difficulty of regularizing the latent space of diffusion models, which we mitigate through the disentanglement constraints on the semantic space.

Results
Semantics-guided diffusion generation framework
Our proposed MolDiffdAE model is an autoencoding framework that utilizes a semantic embedding to control the generation of the equivariant diffusion model. The model has two important features:

• The semantic embedding has maximal mutual information with the input molecule and comprehensively controls all aspects of the generation.

• The disentanglement of the semantic embedding space facilitates the manipulation of multiple objectives, encouraging the model to find directions that achieve all objectives with minimal conflicts.

We implement an encoder that learns the higher-level “semantics” of the input 3D molecule as a continuous, lower-dimensional vector z, aggregating both node- and edge-wise features. The semantic embedding is then provided to an equivariant diffusion model as a condition. The diffusion model functions as a “decoder” that translates the semantic embedding back to the 3D molecular space. (Fig. 1A). During training, the model attempts to predict the clean input data from corrupted data and the semantic embedding. We apply a Wasserstein loss [43] that minimizes the distance between the distribution of the semantic latent space and a prior distribution p(z). In our case, p(z) is an isotropic Gaussian distribution N(0,I). We show that this regularization ensures maximal mutual information between the embedding and the input, and also achieves disentanglement between the dimensions. This leads to our objective function:
L = L (θ) + βL , (1)
dAE D Wass
where L (θ) is the diffusion loss (Eq. 8), L is the Wasserstein loss and β > 0 is the regularization coefficient. The semantic encoder and the diffusion decoder are trained simultaneously on unlabeled 3D molecule conformations from the GEOM dataset [44] in an unsupervised manner. We term this as the pre-trained model.

For the downstream tasks, we could further perform supervised fine-tuning on molecular properties. Specifically, we add an auxiliary classifier on top of the semantic embedding to predict the target property, then train the classifier, encoder and diffusion model in an end-to-end manner with an additional classifier loss:
L = L (θ) + βL + β′L .
dAE-ft D Wass cls
In this work, we only fine-tune on three more complex properties, QED (drug-likeness), SAS (synthetic accessibility) and LogP (lipophilicity, inversely related to water solubility), as we find their prediction with the pre-trained embedding more difficult (Fig. S2). These properties involve intricate interactions between the structural patterns distributed across the atoms, and cannot be learned without supervision, necessitating the fine-tuning. Since property calculation is time-consuming and reliable labeled data is usually rare in real-world scenarios, we construct a fine-tuning dataset by sub-sampling 10% of the pre-training dataset. We denote the pre-trained model as MolDiffdAE and the fine-tuned model as MolDiffdAE-ft.

We then use the diffusion decoder to flexibly perform controlled generation with the semantic embeddings (Fig. 1B, C). As shown in Methods, the embedding z shares maximal mutual information with the input x . Thus, the generated molecule will presumably possess the properties encoded by the embedding. It is also possible to alter certain properties of interest by directly manipulating the embedding. In this way, we transform the challenging problem of controlled generation with diffusion models into a task of latent space optimization (LSO), where several methods could be applied to obtain the desired semantic embedding, which is then translated into the 3D molecular space.

A detailed comparison between molecular generation frameworks can be found in Table 1. Our proposed model combines the high 3D molecular generation capacity of the 3D diffusion model, and the controllability of autoencoding models. In addition, the disentangled semantic space also enables flexible multi-objective manipulation, which is difficult for conditioned generative models that require abundant labeled data and re-training for each objective. These features extend the model capacity beyond less controlled generative tasks such as de novo design and conditioning on single objectives.

Semantics-based molecular retrieval
To interpret the latent semantic space, we perform molecular retrieval from the training set with a query embedding vector z. We use z of existing molecules in the test set as queries. Specifically, we calculate the cosine similarity between each query (source) molecule embedding z and the embedding z* of each molecule in the training set, and then retrieve molecules with the highest cosine similarity to the source.

Shown in Fig. 2A, molecules retrieved by the same query embedding display similar compositions to the source molecule. Quantitatively, the retrieved molecules show significantly higher 2D and 3D Tanimoto similarity to the source molecule, as compared to random pairings (Fig. 2B). In conclusion, we could efficiently query the molecular space and retrieve molecules with desired composition, shape and properties, indicating a semantics-rich latent space that could be used for controlled generation and manipulation.

Probing the property information in the latent semantic space
We train linear regression models to predict the molecular properties from the embeddings, to probe the information content of the semantic space. Despite being trained in an unsupervised manner unaware of any property label, the embedding space of the pre-trained MolDiffdAE could accurately predict many of the properties (Fig. S2A). The dimensions of the embedding can be distinguishably clustered by their differential linear regression coefficients with respect to the properties (Fig. S3). Specifically, only a limited, distinct set of dimensions significantly contribute to each property (Fig. 2C). Also, the embedding space shows strong separation of the high- and low-value groups for many of the properties (Fig. 2E). More complex properties like QED, SAS and LogP are harder to capture with the pre-training. Nevertheless, MolDiffdAE-ft fine-tuned on those properties shows higher predictive performance (Fig. S2B), sparser and more divergent contributions (Fig. 2D), and a better-separated latent space (Fig. 2F), while the fine-tuning has negligible impact on the other properties. Quantitative analysis also shows better disentanglement of the latent space after fine-tuning (Table S2). Meanwhile, the embedding dimensions have minimal internal correlations (Fig. S6A). These results indicate the successful disentanglement between the dimensions enforced by the Wasserstein loss.

Semantics-guided 3D molecule generation
To evaluate the effectiveness of the semantic control over the generation process of the equivariant diffusion model, we explore ways of utilizing it for the generation of 3D molecules with complete semantic information. Specifically, we investigate whether all the molecular properties encoded by the semantics could be successfully translated into the generated molecule.

We obtain the semantic embedding z of known source molecules x from the encoder and perform controlled generation using both deterministic (denoising diffusion implicit model; DDIM [22]) and stochastic (denoising diffusion probabilistic model; DDPM [21]) methods. For DDPM, we generate 10 molecules from each embedding and report metrics on valid molecules only. Fig. S4A shows that DDIM with MolDiffdAE gives near-perfect reconstruction as expected, while DDPM generates molecules that are very close to the source molecule in terms of both composition and 3D shape. This shows that the embedding z dominates the denoising process, while the stochasticity accounts for minor variations.

For a quantitative evaluation of the preservation of the semantics, we define a molecular property profile including 205 non-zero-variance properties calculated by RDKit [50]. We report the following metrics between the generated molecules from DDPM and the source: (1) mean absolute error (MAE) of the properties, aggregated from the selected 2D and 3D property groups (see Table S1 for the full list); (2) Tanimoto similarity (|F ∩F |/|F ∪F | where F ,F are two feature sets) of 2D and 3D fingerprints. We benchmark against the following methods:

• VAE (2D): We use a 2D graph-based variational autoencoder (VAE) to generate random molecules from the embeddings of the source molecules. The generated molecules are represented as SMILES [51] sequences. The 3D properties and 3D Tanimoto similarity of the generated molecules are calculated based on the RDKit-predicted conformations.

• Classifier Guidance: We train a basic equivariant diffusion model without supervision and an independent classifier that predicts all 205 properties. We then use the gradient of the classifier to guide the generation process of a basic diffusion model toward the property profile of the source molecule.

• Random: We use the following methods to obtain the background distribution of the metrics: (1) Property MAE: we randomly sample property values from a N(µ ,σ ) distribution where µ ,σ are the mean and standard deviation of the property, then compare them with those of the source; (2) Tanimoto similarity: for each source molecule, we randomly sample 10 molecules from the training set then calculate their 2D and 3D Tanimoto similarity with the source.

Table 2 shows the aggregated MAEs, along with a detailed breakdown of example individual properties, and average Tanimoto similarities. As expected, the 2D graph-based VAE achieves good preservation of 2D properties, but not of 3D properties. For the 3D structure-based models, MolDiffdAE achieves better preservation of most of the properties than classifier guidance, including both 2D compositional and 3D geometric features. The generated molecules are also more similar to their sources. Notably, unlike classifier guidance where a separate classifier is trained on an external labeled dataset, MolDiffdAE is pre-trained in an unsupervised manner completely unaware of any property. We observe similar performance on randomly sampled semantic embeddings from N(0,I) (Fig. S4B; Supplementary).

These results show that the unsupervised embedding successfully captures the semantic information about the 3D molecular properties and translates it into the generated molecules, indicating effective control over the generation of the equivariant diffusion model. Consequently, leveraging molecular generation based on the semantic embedding of known molecules facilitates the exploration of their neighborhoods in the molecular space in terms of compositions, structures, and properties.

In addition to the generation of full molecules, our model can also adapt to the problem of conformer generation, by fixing the atom and bond types during the denoising process (Supplementary). Despite not being trained specifically on this task, our model shows superior performance to similar models and baselines (Table S5), demonstrating the effectiveness of the semantic control.

Multi-objective molecular property manipulation
In this section, we further explore a generative task that requires more precise control: using the semantic embedding to manipulate the properties of interest, while preserving the majority of composition features, 3D shape and unrelated properties as much as possible. The problem is formulated as follows: let z be the embedding of a molecule x and y = (y ,...,y ) be its property value(s) (p ∈P ,P ⊆P where P is the complete property set and P is the set of properties of interest). Given the desired target property value(s) y′ = (y′ ,...,y′ ), we search the latent space for a new embedding z′ that maps to a new molecule x′ which exhibits the target property y′, while remaining sufficiently similar to x. The latter can be achieved by minimizing the deviation of z′ from z.

We experiment with two ways of obtaining z′:

• Linear manipulation (lm): We first train a linear regression model y = Wz+b (z ∈ R1×d; y ,b ∈ R1×m; W ∈ Rm×d; d is the dimension of z and m is the number of properties) to predict the properties from the embedding, then use it to manipulate the source embedding via:
z′ ← z + W+(y′ − b − Wz) (2)
where W+ ∈ Rd×m is the pseudo-inverse of W, i.e., WW+W = W. It can be proved that z′ is the closest vector to z (in terms of ||z′ − z|| ) that predicts the target property y′ (see Methods and Supplementary for more details).

• Classifier back-propagation (bp): We take the auxiliary classifier of the fine-tuned model and iteratively back-propagate the gradient of the classification loss L to z (Algorithm 1). This is only performed on the fine-tuned properties (QED, SAS, LogP).

On top of direct manipulation, we also experiment with retrieval-augmented generation (RAG) [52] for the fine-tuned properties (QED, SAS, LogP). Specifically, we use the manipulated embedding z′ from classifier back-propagation (bp) to retrieve the 10 closest embeddings from the training set by cosine similarity. The retrieved embeddings (z*) are used to enhance the generation quality via two methods:

• Interpolation (RAG-interp): The retrieved embedding z* is fused with the manipulated embedding z′ through simple interpolation: z′ ← (1−λ)z′ + λz*. We use λ=0.25.

• Second-round manipulation (RAG-reopt): We further manipulate the interpolated embedding toward the same target using classifier back-propagation (bp).

Specifically, we use the deterministic DDIM generation process because it offers stronger control than stochastic DDPM. We first deterministically reverse the diffusion model to map the source molecule x to a noise point x , then perform denoising from x conditioned on the manipulated embedding z′.

We compare our method with the following baselines (see Methods for more details; all diffusion model-based methods start from the deterministically reverse-mapped noise x ):

• VAE (2D)+lm: We perform the same linear manipulation method as MolDiffdAE+lm on the embeddings of the 2D graph VAE. Similar to DDIM, the decoding is also deterministic, as no noise is added to the embedding.

• Conditioned Diffusion: We train a diffusion model conditioned on the property subset P and use it to generate the manipulated molecule conditioned on the target property value y′.

• Classifier Guidance: We train an independent classifier for each property and use it to perform classifier guidance on the basic equivariant diffusion model toward the target property value y′.

• Random Diffusion: We randomly generate molecules with the same atom number as the source using the simple diffusion model with DDPM generation.

• Random MolDiffdAE: We use randomly sampled embeddings from N(0,I) to guide the generation of MolDiffdAE. This reflects the background distribution of the semantic embedding space.

We use a one-pass, deterministic generation method where the encoder and decoder are each called only once per sample. Ideally, the minimal alteration to the semantic embedding during the manipulation implicitly preserves the similarity with the source, and no explicit similarity constraint is needed. For a fair comparison, the same generation mode is also applied to all baselines (see Methods for more details).

On simpler tasks with single properties, our model shows competitive performance with classifier guidance (Fig. S12). We then evaluate our model on more challenging tasks where multiple properties are manipulated simultaneously.

Overall performance on multi-objective manipulation tasks
For the multi-objective tasks, we jointly manipulate 2 or 3 properties selected from halogen count (fr halogen), QED, SAS and LogP. Specifically, we set our objective as lower SAS for easier synthesis, lower LogP for higher water solubility, higher QED for higher drug-likeness, and either higher or lower number of halogens. We define “improvement” for each objective as the change of the value of the property p ∈P adjusted by the desired manipulation direction:
∆ = (yˆ′ − y(0) for increasing objectives (“+”)
pi pi
(y(0) − yˆ′ for decreasing objectives (“–”),
pi pi
where yˆ′ is the property value achieved by the manipulated molecule and y(0) is that of the source molecule. By this definition, more positive improvement is always preferred regardless of the manipulation direction.

We report “success” cases where (1) the generated molecule is valid and novel (not observed in the training set; Fig. S8); (2) the 3D Tanimoto similarity between the generated molecule and the source is greater than 0.5 (95% quantile of random pairwise similarity; Fig. S7); (3) all properties have improvements above a given threshold, where σ is the standard deviation of the property p . We use either a relaxed (>0) or strict (>0.5σ ) improvement threshold. Fig. 3A demonstrates that MolDiffdAE-based methods perform significantly better than the baselines in all but one task, especially under strict improvement thresholds. Furthermore, fine-tuning, back-propagation manipulation and RAG, when applicable, could enhance the overall performance, especially on the more complex properties (QED, SAS, LogP). Meanwhile, the majority of the compositional and structural patterns are well preserved as the manipulated molecules remain similar to their sources qualitatively (Fig. 3C,D,E, I; Fig. S10A). We also show that our methods achieve better numerical preservation of the irrelevant molecular properties, especially with fine-tuning and RAG (Fig. 3B; see Methods for more details).

A detailed analysis of the improvements shows that MolDiffdAE-manipulated molecules generally achieve a better trade-off between all objectives, while classifier guidance occasionally succeeds in one objective but fails in the others (Fig. 3D, F, H; Fig. S10B). These results highlight the advantage of MolDiffdAE’s well-regularized and disentangled latent space, as the objectives have minimal interference with each other. We also observe that the 2D graph-based VAE achieves reasonable success in manipulating the target properties (Table S7), but fails to preserve the 3D similarity, leading to lower final performance. MolDiffdAE, on the other hand, could better preserve both the composition and shape due to the strong semantic control.

In addition, we also analyze the geometric features (Table S4) and conformer energy (Fig. S5A) of the generated molecules, where MolDiffdAE-based methods, especially when combined with classifier back-propagation and RAG, display higher generation quality, indicating a well-regularized latent semantic space that facilitates smooth exploration. In the following sections, we provide a more comprehensive and in-depth analysis of MolDiffdAE’s controlled generation capacity, by investigating its manipulation behavior regarding molecular fragments, preservation of 3D shapes and protein interactions, and ability to manipulate 3D properties.

Case study 1: Joint manipulation of LogP and SAS
We perform an in-depth comparison between the manipulation results targeting both lower LogP (higher water solubility) and lower SAS (higher synthetic accessibility) for MolDiffdAE-based manipulation and classifier guidance. These properties are ideal targets for evaluating the impact of the disentangled latent space, as both of them are defined by the combined contribution of the molecular fragments: LogP is calculated as the sum of the hydrophilicity score of each atom [53], and SAS based on the frequencies of the fragments, as more common fragments imply easier synthesis [54]. As a result, different fragments may have aligned or diverging impacts on LogP and SAS. Thus, it is important for the property manipulation process to separate their differential contributions, and introduce fragments that contribute to both objectives.

In the first example (Fig. 4A), classifier guidance successfully reduces the SAS score by replacing the low accessibility backbone with a high accessibility one. However, some of these newly introduced fragments also reduce water solubility. For example, additional benzene rings are highly accessible but less water-soluble. Similarly, in Fig. 4B, the low accessibility backbone is modified, but solubility is not improved as the new backbone does not include more hydrophilic atoms. In comparison, manipulation with MolDiffdAE successfully considers both objectives, as it introduces fragments that are both more synthetically accessible and more water-soluble.

We further complicate the task by adding “more halogens” as another objective. This objective could interfere with the reduction of LogP as halogens are hydrophobic. Meanwhile, aromatic carbons bound to halogens would slightly increase water solubility as a result of the polarization of the aromatic ring. We show that on top of the aforementioned behaviors, MolDiffdAE-based manipulation also tends to attach halogens to aromatic rings to counteract the hydrophobic effect of halogens, while avoiding aliphatic carbons (Fig. 4C, D). In contrast, classifier guidance could trivially succeed in increasing the number of halogens, but fail to make other necessary modifications accordingly to reduce LogP (Fig. 4C) or SAS (Fig. 4D).

These results further prove that the disentangled semantic space successfully guides the diffusion model in multi-objective manipulation tasks. Manipulation in the disentangled semantic embedding space of MolDiffdAE could identify directions that simultaneously satisfy all objectives. With simple classifier guidance, however, the gradient of the classifiers could lead to conflicting directions and interfere with each other without proper regularization.

Case study 2: Validity enhancement by retrieval-augmented generation
We observe that using RAG achieves competitive results with MolDiffdAE-ft+bp, the best-performing method without RAG, in the tasks involving QED, SAS and LogP, and noticeably better results under relaxed improvement thresholds (> 0). A detailed analysis shows comparable performance in terms of success rate (Fig. 3A) and similarity to the source (Fig. 5A, top) between the MolDiffdAE-based methods, while RAG noticeably improves the validity of the generated molecules (Fig. 5A, middle). We also show that the generation quality is overall higher with RAG, in terms of the distribution of geometric properties (Fig. 5A, bottom; Table S4). We notice no significant performance differences between the two RAG methods tested.

Fig. 5B shows that in some of the success cases, the retrieved molecules already contain the desired target properties, due to an information-rich and structured latent space. More importantly, the generated molecule incorporates compositional features from both the source molecule and the retrieved molecule, potentially enhancing the generation quality by keeping the manipulated embedding closer to the data distribution.

External Application 1: Binding-preserving property manipulation
Due to its ability to preserve 3D shapes, which is closely related to the interaction conformations, our semantics-guided property manipulation is particularly useful for protein-binding ligands. In those cases, the binding toward the target protein should be maintained, while some useful properties such as QED, LogP and SAS are modified. To test the generalizability of our method to this task, we apply both our pre-trained and fine-tuned model to the external test set of CrossDocked [55], which contains the 3D structure of 100 ligand-receptor pairs. We observe that MolDiffdAE-based methods show competitive performance with classifier guidance on single-objective tasks (Table S6), and significantly higher performance on multi-objective tasks (Fig. 6A).

We then dock the manipulated molecules back to their respective receptors using QuickVina [56, 57]. Fig. 6B shows examples where the docking is mostly maintained after manipulating the target properties, as reflected by similar docking scores. Overall, on top of successes in improving the target properties, MolDiffdAE-based manipulation could also preserve the Vina score distribution of the source molecules. Meanwhile, classifier guidance preserves the 3D shape and docking but fails in the manipulation objective (Fig. 6A), while 2D VAE-based manipulation breaks the docking (Fig. 6C, D). These results are particularly compelling as our MolDiffdAE model is not explicitly conditioned on the structure of the receptor, but could still maintain the binding throughout the property manipulation. Thus, it could be used to manipulate lead compounds without the need to jointly model both the ligand and the receptor, as the preservation of the 3D shape implicitly preserves the docking. As reliable structures for ligand-receptor complexes are relatively rare, our method could efficiently transfer the 3D shape knowledge from the more data-rich domain of isolated 3D molecules to the generation and manipulation of ligands. It is noteworthy that the method only preserves the 3D shape not necessarily the binding itself, which would require fine-tuning with the protein pocket as the direct conditioning. We leave these for further studies.

External Application 2: Manipulation of DRD2 binding affinity
Beyond basic molecular properties, we further test our method on more complex external properties. We focus on increasing the binding affinity to the dopamine D2 receptor (DRD2). Despite being a well-studied benchmark, the 3D structural data for positive DRD2 binders is still relatively scarce. We select 1400 positive structures from GEOM using the oracle predictor provided by the Therapeutics Data Commons (TDC) [58]. Due to the complex nature of protein binding, the small data size makes it difficult to train reliable 3D structure-based classifiers or conditioned diffusion models. We fine-tune our MolDiffdAE model on the DRD2 binding dataset, achieving good separation of positives and negatives in the semantic space (Fig. 7A), and perform linear manipulation to increase the binding affinity score. The generated molecules are then evaluated by the same oracle model, and we report positive molecules with binding score >0.5.

As shown in Fig. 7B, MolDiffdAE-ft with linear manipulation outperforms all baselines with a simple one-pass generation and no post-generation filtering. This implies that fine-tuning of the latent semantic space enables efficient use of the limited labeled data. We observe that the manipulated molecules have more hydrophobic components, such as aliphatic carbon (Fig. 7C), naturally resulting in higher LogP (Fig. 7D, bottom). This encourages binding toward DRD2’s highly hydrophobic pocket. Also, though not an explicit objective, the manipulated molecules also show higher QED (Fig. 7D, top), indicating underlying mutual factors that contribute to both DRD2 binding and drug-likeness. However, the manipulation also increases synthetic difficulty (higher SAS; Fig. 7D, middle). This could be potentially improved by jointly manipulating SAS as well, as in previous experiments.

Manipulation of 3D properties
Most of the existing works for controlled molecular generation are restricted to compositional properties, due to their limitations of relying on sequence or 2D graph representations, while the manipulation of 3D geometric properties remains poorly explored. MolDiffdAE as an equivariant 3D diffusion model enables operation on those properties. We manipulate asphericity (Asp) and radius of gyration (RoG) with linear manipulation. As expected, the 2D VAE fails in controlling the 3D properties. Among the 3D equivariant model-based methods, MolDiffdAE shows superior performance in both single- and multi-objective tasks, especially under more stringent improvement thresholds (>0.5σ ) (Fig. 8A).

Contrary to our previous results on compositional properties where the composition undergoes drastic changes while the 3D shape is mostly preserved, manipulation of the 3D properties typically leads to major alterations in the 3D shapes compared to the source molecules (Fig. 8B). In agreement with common observations, molecules with higher asphericity are less spherical, while molecules with higher radius of gyration are more spatially stretched. Meanwhile, the composition of the source is well-preserved after the manipulation. We also observe a good trade-off between the objectives with MolDiffdAE, while classifier guidance either fails in achieving both objectives, or generates few valid molecules (Fig. 8C). The latter suggests the complexity of 3D geometries and implies the usefulness of encoding the 3D structural information into the molecular semantics. Such tasks would be useful in the search of candidate molecules with a preferred 3D shape and a fixed chemical composition.

Discussion
In this study, we present MolDiffdAE, an equivariant diffusion model guided by a disentangled semantic latent space for the controlled generation of 3D molecules. We show that the semantic embedding effectively controls the composition, shape, and molecular properties of the generated molecules, in both stochastic generation and deterministic property manipulation. More importantly, the regularization-induced disentanglement of the semantic embedding empowers effective multi-objective manipulation, which is a challenge for traditional diffusion models where the latent space is entangled and difficult to operate on. Though trained in an unsupervised manner, our pre-trained model is capable of capturing simple compositional and 3D shape properties, and can be used to manipulate them through the semantic embedding. For more complex properties whose contributing factors are distributed across many atoms, such as drug-likeness, solubility, synthetic accessibility and protein binding affinity, we show that our model could be further fine-tuned on small, external property-labeled datasets to enhance the performance.

Several existing controlled molecule generation frameworks on 2D molecules deploy iterative optimization or post-generation filtering to ensure the generation quality and the similarity to the source molecule if required [1, 29]. Our method, however, uses one-pass generation, calling the diffusion decoder only once, and evaluating only the final generation output. Similarity constraint is implicitly enforced by minimizing the distance to the source semantic embedding. Despite the simplified generation pipeline, our method already achieves a good success rate with respect to molecular validity, property values and similarity to the source, as a result of the strong guidance of the semantic embedding and the high capacity of the diffusion model. We also show that the information-rich, low-dimensional semantic embedding vector can be efficiently used for the retrieval of similar molecules and subsequently the enhancement of manipulation results through RAG. Nevertheless, the generation quality of our model could also be further improved through iterative refinement or encoding of the molecular similarity into the latent semantic space with contrastive learning. It is also worth noting that in our experiments, we use randomly sampled “negative” source molecules, some of which may be too challenging for the target value. In real applications, a rational pre-selection of more promising source molecules could further enhance the success rate.

Beyond basic molecular properties, we expect the method to be applicable to a wide spectrum of objectives, such as toxicity, conformational energy, multi-receptor binding affinity, and material properties. It is also possible to associate the semantic embedding with text descriptions, and utilize the recent advances in large language models (LLM) for more flexible conditioning and control [32]. Our framework could also flexibly incorporate recent advances in representation learning and diffusion model architectures for 3D molecules, to enhance the representativeness of the latent space, the quality of generation, and the generalization to rarer and more complex structures. Specifically, using generative models with greater flexibility could enable more drastic manipulations, which is common in real-world drug design and optimization.

Furthermore, although we have demonstrated the preservation of the 3D shape and the binding pose for ligand molecules independent of the receptor proteins, the semantics-controlled generation of ligands could still further benefit from jointly modeling both the molecular ligand and the receptor protein given sufficient data. Our framework could also be potentially incorporated into diffusion models for pocket-aware ligand design [34, 59], by manipulating the ligand within the confinement of the pocket structure.

In summary, we present a comprehensive exploration of controlled generation of 3D molecules using diffusion models. The design of our framework encourages several desirable behaviors for molecule design, such as precise control over the generative process, disentanglement of the semantic space and flexibility in adapting to diverse tasks requiring control of molecular properties. We demonstrate superior performance in the more challenging tasks with multiple properties and limited data compared to existing baselines, which highlights its potential in extending to real-world applications. We believe that our model, supported by rigorous evaluation and reproducibility, provides a solid foundation for future experimental studies, and offers new hypotheses for further investigation. The initial results of our approach are promising and may inform and guide subsequent validation efforts.

Methods
Data
Datasets
Pre-training dataset: For the pre-training of the MolDiffdAE model, we use the drug-like molecules from GEOM [44] with the same training/validation/test split as previous works [25] as the pre-training dataset.

Fine-tuning dataset (RDKit property): We also build a property dataset for the fine-tuning of MolDiffdAE and training of the property classifiers. Specifically, we use RDKit (version 2024.3.2) to calculate the properties of the samples in the validation set of the pre-training dataset, resulting in a dataset comprising 10% of the pre-training set.

Fine-tuning dataset (DRD2 binding): We use the oracle model from the Therapeutics Data Commons (TDC) [58] (version 0.4.1) to calculate the DRD2 binding affinity of all molecules from GEOM. We merge the 1400 positive structures (predicted affinity score > 0.5) from the training and test set, along with 5× randomly selected negatives, for the fine-tuning.

CrossDocked test set: We only use the 100-sample test set following the same split as previous works [59, 60] for the property manipulation. The molecular properties of the ligands are calculated by RDKit.

A detailed pipeline of the construction of the pre-training and fine-tuning datasets can be found in Fig. S1. For property manipulation tasks, we use the same fine-tuning dataset for the training of the linear (for linear manipulation) or EGNN-based (for classifier guidance) classifiers. All generation and manipulation results are reported on the test set.

Representation of the 3D molecular data
For simplicity, we remove hydrogens from the molecules. We represent each 3D molecule as a fully-connected 3D graph x = (r,h,E) where r is the atom coordinates, h is the atom type and E is the edge type. Following previous practices on equivariant diffusion models [15], r is centered on the center of mass (CoM) of the molecule. h and E are represented by one-hot encoding. Aromatic bonds are considered a distinct type beside single, double and triple bonds.

Semantics-guided autoencoding equivariant diffusion model
Semantics encoder
Let x = (r ,h ,E ) be an input 3D molecule. We design an encoder with an equivariant backbone that learns the conditional distribution of the semantic embedding z given the input:
q(z|x ) = N(µ ,σ ), µ ,σ = Encoder (x ) (4)
0 z z z z γ 0
where γ is the encoder parameter. z is then sampled from the distribution and provided to a diffusion “decoder” to generate a reconstruction of the input. The embedding z is treated as a condition of the diffusion process. In practice, we use a deterministic encoder where only µ is calculated and σ = 0.
z z

In addition, the encoder can be co-trained with an auxiliary classifier ψ that predicts molecular properties y of interest from z, encouraging z to also carry information about y. We use mean squared error (MSE) as the auxiliary classification loss L for regression tasks (RDKit properties) and binary cross-entropy (BCE) for binary classification tasks (DRD2).

Semantics-guided diffusion decoder
The diffusion decoder is based on the denoising diffusion probabilistic model (DDPM) [21] with the semantic embedding as the condition. The conditional data distribution p (x |z) is approximated through a series of latent variables x ,...,x with the same dimension as x , called the reverse process, starting from a random noise point x , parametrized by θ:
p (x |z) = p (x |z)dx = p(x ) p (x |x ,z)dx . (5)
θ 0 θ 0:T 1:T T θ t−1 t 1:T
t=1

The forward process q(x |x ) gradually adds noise to x until it eventually becomes random noise x :
q(x |x ) = q(x |x ). (6)
1:T 0 t t−1
t=1

The objective is to maximize the ELBO of log p(x ) = log p (x |z)p(z)dz. Under Gaussian assumptions, it is equivalent to minimizing the prediction loss of either the clean data x , or the noise ϵ added to x at time step t, from the corrupted data x [21]. Typically, the noise parametrization is employed to achieve greater training stability. However, previous studies on 3D molecules [24, 25, 61] show the advantage of using the clean data parametrization, as the model is aware of the overall graph structure throughout the training.

Based on these considerations, we train the diffusion model to predict the clean input x = (r ,h ,E ) at time step t from corrupted data x = (r ,h ,E ), conditioned on the embedding z:
rˆ ,hˆ ,Eˆ = DM (r ,h ,E ,z,t). (7)
0,t 0,t 0,t θ t t t
Thus, the diffusion objective is defined as:
L (θ) = E [||rˆ − r ||2 + ||hˆ − h ||2 + ||Eˆ − E ||2]. (8)
D (r0,h0,E0),(rt,ht,Et) 0,t 0 2 0,t 0 2 0,t 0 2
t=1

Intuitively, z can be considered as controlling the “direction” of the denoising (i.e., generation) toward the desired semantics. In Methods (“Regularization of the semantic embedding” and “Theoretical perspective of the objective”), we describe a regularization term to enforce maximal mutual information (MI) between z and the input, which empowers z to effectively guide and control the generation process.

Regularization of the semantic embedding
To control the scale and shape of the distribution of z, we employ a Wasserstein loss [43] on the conditional distribution q(z) = q(z|x )q(x )dx . Specifically, we use a sample-based kernel maximum mean discrepancy (MMD) [62] on mini-batches of size n to make q(z) approach the shape of a Gaussian prior p(z) = N(0,I). We use the linear time unbiased estimator [62, 63]:
L := MMD2(q(z),p(z)) = h(z ,z˜ ,z ,z˜ ) (9)
Wass ⌊n/2⌋ 2i−1 2i−1 2i 2i
i=1
h(z ,z˜ ,z ,z˜ ) = k(z ,z ) + k(z˜ ,z˜ ) − k(z ,z˜ ) − k(z ,z˜ ),
i i j j i j i j i j j i
where k is the kernel function, z are encoded from the data points in the minibatch and z˜ are randomly sampled from p(z)=N(0,I).

Theoretical perspective of the objective
One concern with the autoencoding diffusion model framework is that the decoder may ignore the embedding z and solely rely on the noise x for the generation. This mirrors a similar problem observed in several generative models [64, 65]. We show that our objective in Eq. 1 also ensures maximal mutual information between z and the input x , which leads to a more informative latent space.

When β = 1 (proof in Supplementary):
L = L (θ) + KL(q(z)||p(z)) = −L − MI(x ,z) (10)
D ELBO 0
where L := −L (θ) − E KL(q(z|x )||p(z)). (11)
ELBO D q(x0) 0
KL stands for the Kullback-Leibler divergence. L is the ELBO for log p (x ) under our model assumption. L (θ) is the same objective as Eq. 8. This means that minimizing L encourages a solution with both high ELBO and high mutual information.

In practice, KL(q(z)||p(z)) is generally intractable. We therefore use the MMD form L (Eq. 9) as a surrogate, which still matches q(z) and p(z) while empirically preserving semantic information in z. This leads to our objective function Eq. 1.

Diffusion generation process
In this study, we experiment with different generative methods for diffusion models. For semantics-guided stochastic generation, we use DDPM [21] which progressively samples the clean data from random noise. For property manipulation tasks, we employ the deterministic DDIM sampling [22]. DDIM is based on a non-Markovian forward process, enabling deterministic generation, while trained with the same objective as DDPM. Starting from the random noise, we could obtain a reconstruction of the clean data by progressively removing the predicted noise deterministically:
v = √ α v t − √ 1−α t ϵˆ( t v) + √ 1−α ϵˆ(v), ϵˆ(v) = √ 1 v − √ α t vˆ , (12)
t−1 t−1 α t−1 t t 1−α t 1−α 0,t
t t t
where v ∈{r,h,E}, v is the corrupted value and vˆ is the diffusion model prediction of the clean data from x and z. This could also be rewritten as:
v = √ α vˆ − √ 1−α √ t−1( √ α vˆ − v ). (13)
t−1 t−1 0,t t 0,t t
1−α
t
When the time step is sufficiently small, we could also map the input data (r ,h ,E ) to the noise point (r ,h ,E ) through an inverse of the denoising process [22, 37]:
v = √ α vˆ − √ 1−α √ t ( √ α vˆ − v ). (14)
t t 0,t−1 1−α t−1 0,t−1 t−1
t−1
When z remains unchanged in both processes, we could achieve perfect reconstruction like in typical autoencoders. When z is altered, the reverse process would be steered to a different direction than reconstruction.

Model architecture
Satorras et al. [66] proposed the E(n)-equivariant graph neural network (EGNN), aiming to incorporate geometric symmetry into molecular modeling. Given a fully connected graph where each node is represented as a coordinate r ∈ R3 and a feature vector h ∈ Rd, and a function (r′,h′)=f(r,h), the EGNN ensures the equivariance constraint. Namely, given an orthogonal matrix R ∈ R3×3 and a translation vector t ∈ R3:
T(Rr′+t,h′) = f(Rr+t,h). (15)
i.e., h is invariant to transformations while r receives an equivalent transformation.

Hoogeboom et al. [15] use the EGNN as the backbone of a DDPM for 3D molecules. It is shown that with an invariant prior p(x ), the equivariant process results in an invariant data distribution p (X), where the probability of a data point remains the same after transformation [67]. This greatly improves data efficiency and model generalizability.

We use the EGNN backbone from MolDiff [25] for both the encoder and the diffusion model. Given the input (r,h,E) and some external input vector u, the values on node i and edge (i,j) are updated with the following message passing process:
E˜ = ϕ (E ,||r − r || )
ij d ij i j 2
h = ω (h ) + ϕ (h ,E˜ ,u)
i h i h j ij
j
E = ϕ (h ,E˜ ,u) + ϕ (h ,E˜ ,u) + ω (h ) + ω (h ) + ω (E˜ )
ij e k ki e k kj h i h j e ij
k k
r = r + ϕ (h ,h ,E˜ ,u) r i − r j , (16)
i i r i j ij ||r − r ||2
j i j 2
where ω ,ω are linear embedding layers; ϕ ,ϕ ,ϕ ,ϕ are multi-layer perceptrons (MLPs). The vector input u depends on the model: for simple diffusion models, it is the time step t; for the diffusion decoder of MolDiffdAE, it is the concatenation of the semantic embedding z and t. The semantic encoder takes no vector input.

For the semantic encoder, we aggregate the invariant output of the EGNN with a MLP:
z(r),z(h),z(E) = EGNN (r ,h ,E ),
γ 0 0 0
z = MLP(z(h),z(E)), (17)
where z(r) is equivariant and z(h),z(E) are invariant. Because all properties we investigate in this study are only related to the invariant features (atom types, edge types and pairwise distances), we only use the invariant output of the encoder to calculate the semantic embedding.

The diffusion model uses the same EGNN architecture to predict the clean data x from the noisy data x as in Eq. 7. The semantic embedding z and the embedding of the time step t are concatenated to all node and edge features.
rˆ ,hˆ ,Eˆ = EGNN (r ,h ,E ,z,t). (18)
0,t 0,t 0,t θ t t t
Since the DDIM-based diffusion inversion (Eq. 14) only works on continuous distributions [22], we use a continuous diffusion model for our experiments. The categorical values h and E are treated as multivariate Gaussian distributions with the respective dimensions. Details of the model hyperparameters can be found in Supplementary.

Semantics-guided controlled generation of 3D molecules
Given the semantic embedding z of a known molecule x = (r ,h ,E ) from the encoder, we combine it with randomly sampled stochastic noise x′ = (r′ ,h′ ,E′ ) to generate new molecules x′ = (r′,h′,E′).

Specifically, we obtain the embedding z of 100 source molecules randomly selected from the test set and generate 10 molecules controlled by each z. For each generated molecule x′, we use RDKit to obtain its molecular properties and calculate the mean absolute error (MAE) for each property compared to that of the source x . We report the MAEs for individual properties as well as property sets (“All”, “All 2D” and “All 3D”; the full list of properties in each subset can be found in Table S1). For the latter, we discard the properties with standard deviation < 0.1 and perform z-scoring (subtracting the mean and scaling by the standard deviation) on the remaining properties. As a random control, we calculate the MAE of 10 random property profiles for each source molecule, where the value of each property p is sampled from N(µ ,σ ) where µ ,σ are the mean and standard deviation.

We also report the average 2D and 3D Tanimoto similarity between the generated molecules and their sources (see Methods). The random control is the average similarity between 10 randomly sampled molecules from the training set and each source molecule.

Multi-objective property manipulation methods
Linear manipulation
Given a known molecule, we use a simple linear model-based framework to manipulate its semantic embedding to gain certain properties. We train a linear regression model on the training set to predict the property(ies) from the semantic embedding. Then, given a source embedding z, target value(s) y′, and the weight and bias of the linear regression (W,b), we aim to obtain a new embedding z′ that predicts y′ by the linear model while being sufficiently close to z. This can be formulated as a constrained optimization problem:
min ||z′ − z||2 subject to y′ = Wz′ + b. (19)
z′ 2 p
This problem has a closed form solution (Proof in Supplementary):
z′ = z + W+(y′ − b − Wz). (20)
The pseudoinverse W+ is calculated using singular value decomposition [68].

Classifier back-propagation manipulation
For the fine-tuning properties (QED, SAS, LogP), we could also use the auxiliary classifier ψ to directly manipulate the embedding through iterative back-propagation (Algorithm 1).

Algorithm 1 Classifier back-propagation manipulation
Embedding z, target property value y′, learning rate λ, number of iterations n
Auxiliary classifier ψ
z′ ← z
for i in [0,n] do
yˆ = ψ (z′)
L = MSE(yˆ,y′)
cls
z′ ← z′ − λ∇ L
z′ cls
end for
We choose n=20, λ=1.0 for QED and λ=0.1 for SAS and LogP.

Retrieval-augmented generation (RAG)
We calculate and store the semantic embeddings for all molecules in the training set. Given a manipulated embedding z′ obtained through any means, we query the training set and retrieve the top 10 samples whose embeddings have the highest cosine similarity to z′. z′ is then fused with each retrieved embedding z* through linear interpolation:
z′ ← (1−λ)z′ + λz* . (21)
We experiment with λ = 0.1,0.25,0.5. We could also further optimize the fused embedding with the same methods described above.

Deterministic generation from manipulated embeddings
To decode the manipulated embedding, we first reverse map x to a noisy data point x with the source embedding z. We then provide the manipulated embedding z′ and x to the diffusion decoder to generate a manipulated molecule x′ = (r′,h′,E′). In practice, we only reverse map to x instead of x and then denoise for T/2 steps without skipping, as it gives the best generation quality empirically.

For each target property set P , we report results on source molecules from the test set whose value of each property p ∈ P is above (decreasing objectives) or below (increasing objectives) the mean µ . We select 100 source molecules for the RDKit properties and 200 for DRD2 binding. For CrossDocked samples, we evaluate on all the 100 test samples. To shift the property value toward the desired direction, the target value is set to y′ = µ + 2σ (increasing objectives) or y′ = µ − 2σ (decreasing objectives) for the RDKit properties, and y′ = 1 for DRD2 binding.

Baselines
VAE
We use the 2D graph-based VAE model HierVAE [28] pre-trained on the ChEMBL dataset [69]. Input molecules are represented as SMILES sequences. Given the source molecule, we use the HierVAE encoder to obtain its embedding. We add random noise to the embedding for semantics-guided generation, but not for property manipulation. For property manipulation tasks, we train a linear regression on the embeddings and manipulate the embedding in the same way as in previous property manipulation tasks. The corrupted or manipulated embedding is then provided to the HierVAE decoder to obtain the resulting molecule.

For a fair comparison with the unsupervised pre-training and one-pass deterministic generation mode of MolDiffdAE, we adopt a similar setting for HierVAE. Specifically, we use the HierVAE model pre-trained on isolated data, instead of the specialized graph-to-graph model set up for graph translation and trained on paired data.

Simple diffusion model
For the simple diffusion model, we use the same model architecture and hyperparameters provided by MolDiff [25] (https://github.com/pengxingang/MolDiff). As our MolDiffdAE model is based on continuous diffusion, we also use the continuous version of MolDiff.

Conditioned diffusion model
For each property or property set P , we train a diffusion model on the property dataset conditioned on the property(ies) y , i.e., the reverse and forward processes are defined as:
p (x |y ) = p(x ) p (x |x ,y )dx . (22)
θ 0 p T θ t−1 t p 1:T
t=1
q(x |x ,y ) = q(x |x ,y ). (23)
1:T 0 p t t−1 p
t=1
In practice, the property value is combined with the time step t and then concatenated to every node and edge feature. The model is trained in the same way as the simple diffusion model baseline.

For the DDIM-based property manipulation, we first reverse map x to a noisy data point x with the source property y . Similar to the property manipulation task, we perform reverse mapping with T/2 steps. Then, x is combined with the target value y′ in the denoising process to obtain the manipulated molecule x′. For multi-objective tasks, we train one model conditioned on all properties of interest (Halogen, QED, LogP and SAS for the 2D property model; asphericity and radius of gyration for the 3D property model). During generation, we use the target value for the manipulated dimensions of the property vector and keep the source’s original value on the other dimensions.

Classifier guidance
Classifier guidance is performed on the continuous version of MolDiff. For property-controlled generation, we train classifiers that predict all RDKit properties. For property manipulation, we train separate classifiers for each target property on the property set. The classifiers have the same EGNN-based architecture as MolDiffdAE’s semantic encoder, with a 2-layer MLP prediction head added on top of the model output. During training, the input is corrupted with random noise sampled from the same noise schedule as the diffusion model.

We use DDPM for property-controlled generation and DDIM with T/2 reverse and denoising steps for property manipulation. In either case, let x = (r ,h ,E ) be the prediction output of any denoising step, P be the set of target properties, y′ be the target value of property p and Ψ be the respective classifier, at every time step, each modality of x = (r ,h ,E ) is further updated with the classifier gradient:
v = v + λ ∇ L (Ψ (x ),y′ ), (24)
t t pi vt cls pi t pi
pi∈Psub
where v ∈{r,h,E}, λ is the guidance strength and L is the classifier loss, which is MSE in our experiments. For property-controlled generation, we perform min-max scaling on all properties and use λ = 0.1. For property manipulation, we use λ = 0.01 for SAS and LogP, λ = 0.1 for the other properties.

Evaluation of generated molecules
Validity, similarity and novelty
To evaluate the quality of generated molecules, we use the following metrics:

• Validity Valid molecules are defined as those that have no violation of valency or disconnected atoms. In practice, this is evaluated by whether the molecule could be successfully processed by RDKit’s SanitizeMol and Kekulize functions.

• Novelty We compare the SMILES of the generated molecules to the training set and report molecules with no overlap as novel. In our experiments, almost all generated molecules are novel.

• Similarity We use RDKit’s TanimotoSimilarity function with default parameters and RDKFingerprint to calculate the 2D Tanimoto similarity, and the ShapeTanimotoDist function to calculate the 3D Tanimoto similarity (ShapeTanimotoSimilarity = 1−ShapeTanimotoDist). In addition, for 3D similarity, the molecules are aligned using GetCrippenO3A from the rdMolAlign module.

Molecular property
The property values of the generated molecules are calculated with RDKit [50] and the TDC [58] oracle, respectively, in the same manner as the fine-tuning dataset.

To evaluate whether the “irrelevant” properties are modified during the manipulation, we select a list of 2D compositional and 3D shape properties that do not overlap with the target properties (all properties in Table S1 excluding fr halogen). For all generated samples, we calculate the mean absolute error (MAE) of the selected irrelevant properties, normalized by the standard deviation of each property.

The contributions of atoms and fragments to properties like LogP and SAS are retrieved from the source code. Specifically, for the labeling in Fig. 4, “high accessibility” fragments are those whose raw contribution score is >2 while “low accessibility” <−0.5. It’s worth noting that RDKit’s SAS output inverts the sum of the contribution scores, so higher contribution leads to lower SAS (i.e., higher synthetic accessibility). For LogP, the labels are based on the scoring table in [53], where negative scores make the molecule more soluble.

Docking
For the CrossDocked test samples, the generated molecules are docked to the respective receptor using QuickVina (version 1.2.5) [56, 57]. Specifically, the coordinates of the generated molecules are shifted to the center point of the source molecule in the original docking conformation. The density function of docking scores is estimated using Gaussian kernel density estimation [70].

Peer review rounds
[{"round": 1, "reviewer_comments": {"reviewer_1": "The authors present an approach to conditionally generate or “edit” molecules in 3D using an equivariant diffusion model\nwith a disentangled autoencoder. This approach enables the editing of some molecule properties whilst preserving others -\na valuable concept for computational drug design. In particular, they point out that this approach is better than alternatives for\nmulti-property optimization (testing the change of up to 3 properties). Please see my comments regarding the paper.\n1. Please include a validity analysis. It is stated that you only consider valid molecules but I do not see any evaluation of\nhow many molecules are valid for those generated. Moreover, the geometrical validity (e.g., bond lengths, bond angle\ndistributions etc.) should be evaluated so that readers can understand how geometrically valid molecules are. This is\ntypically evaluated by other models for 3D molecule generation. For example, you could use Validity3D, PoseCheck,\nPoseBusters and other such tools to measure this.\n2. Table 1 is not a “detailed comparison” as stated by the authors. And lacks sufficient consideration of alternative methods\nfor generation in 3D. There are many more methods than a 2D autoencoder and 3D diffusion. Please include a balanced\nand broader comparison in Table 1.\n3. In the main text, qualitative conclusions are drawn for the fine-tuned model between Figure 2A left and right. This is\nspeculative at best and to make such conclusions qualitative results should be included in the text. Additionally, the\ncorresponding Figure S2 should include statistics so that improvements in prediction are quantifiable.\nSeveral points regarding Section 2.4.3,\n4. The molecules in Figure 7 are illegible and should be represented in 3D as is done elsewhere.\n5. Figure 7 is also incorrectly cited as Figure S7 twice in the main text.\n6. The last sentence concluding that this approach enables modelling 3D ligands without jointly modelling the protein could\nbe a moot point. This is because the 3D bioactive conformation will still need to be known in the context of binding\npreserving, and therefore the complex structure will still need to be known. Otherwise you risk modifying 2D topology while\npreserving an incorrect 3D structure which is irrelevant, without guarantee that the new 2D topology can form the correct 3D\nconformation.\n8. More important in binding preservation is the preservation of key protein-ligand interactions. Only considering the docking\nscore may slightly overlook this problem. Qualitatively I can see examples in Figure 7 where key interactions look to be lost.\nThis at least warrants a discussion, for example, perhaps modelling polar hydrogens will improve this.\n7. Is it possible to conduct experimental evaluation of some examples to ensure that binding is in fact preserved? I would\nexpect this level of scientific rigor in Nature Communications.\n8. The CrossDocked models are not “generally out of distribution” according to Figure S7(a). They overlap in chemical\nspace and have distribution subsets that overlap with GEOM. It would be more robust to do a pairwise comparison of the 100\nCrossDocked test molecules to find the closest molecules in GEOM drugs using Tanimoto similarity.\n9. Another potentially useful application of this approach is to preserve both the 2D and 3D molecule while generating\nalternative conformations. Is this possible with the proposed approach? This could also be an effective way to measure the\nscope of the model.\n10. Please provide specific information on the oracle used in Section 2.4.4 (versions, weblinks etc.).\n11. I don’t understand the axis of Figure S3, please add labels. I assume the x axis is (cid:1), but are there properties not labelled\non the y axis. Why are there gaps?\n12. Please include details on training and inference time.\n13. Other diffusion models have been proposed for multi-parameter optimization such as PILOT by Cremer et al. that haven’t\nbeen referenced or discussed.\nCremer, Julian, et al. \"PILOT: equivariant diffusion for pocket-conditioned de novo ligand generation with multi-objective\nguidance via importance sampling.\" Chemical Science 15.36 (2024): 14954-14967.\n(Remarks on code availability)\nThe README does not contain sufficiently detailed instructions for installation and running of the code.", "reviewer_2": "The article by Li et al. presents MolDiffdAE which is a data-driven model for the conditional generation of 3D molecular\nstructures. The article is technical and generally well-written, however, I do not believe it is suitable for the broader\nreadership of nat. commun. Although, I enjoyed the discussion and Authors’ of “semantics” and, in particular, controlled\nmulti-objective optimization, the article is mostly a report of incremental technical advancements with functionality\ndemonstrated using (mostly) typical benchmark properties that belongs in a more specialized journal. The field of 3D\nmolecular generation has steadily grown (for example, Peng et al., \"MolDiff: Addressing the Atom-Bond Inconsistency\nProblem in 3D Molecule Diffusion Generation,\" ICML 2023. (please fix ref 25), references therein and thereof, to name a\nfew), and, therefore, I believe that MolDiffdAE does not possess the broader applicability or impact that is expect of a Nature\nCommunication article.\n(Remarks on code availability)\nThe code is well organized and has a high probability of resulting in reproducibility. I also appreciate that the Authors have\nincluded relevant references in the code as comments.\nAt first glance, there are a few places where pep8 style guidelines are not followed. This is not essential to resolve, but it\ncould improve readability.", "reviewer_3": "The manuscript “Disentangled Autoencoding Equivariant Diffusion Model for Controlled Generation of 3D Molecules”\npresents a method that combines a diffusion model with semantics learned from an autoencoder to control the generation of\n3D molecules with desired properties. The authors state that this model can support control over the properties of the\ngenerated molecules. According to the results, the method provides strong performance and surpasses the baselines.\nOverall, the manuscript aligns with the scope of nature communication, and the proposed method appears physically and\nmathematically sound, contributing to advancements in the generative models for molecules. However, several concerns\nneed to be addressed to clarify this manuscript before publication.\n1. Questions related to methodology\n• This manuscript discusses the model's ability to manipulate molecular properties through disentanglement in the latent\nspace, particularly in multi-property scenarios. If accurate, the method should modify selected properties along user-guided\ndirections while keeping other properties largely unchanged due to the disentangled latent dimensions. The manuscript\nmentions the modification of selected properties but does not address the impact on non-selected properties.\n• On page 11, the authors wrote:\nOn simpler tasks with single properties, our model shows competitive performance with classifier guidance (Supplementary\n2.3; Fig. S6). We then evaluate our model on more challenging tasks where multiple properties are manipulated\nsimultaneously.\nIn Fig. S6, the model appears to modify some properties as directed. However, if the model can modify a single property as\ndirected, it should also maintain other properties relatively unchanged. In Fig. S6, the results show modifications to QED,\nLogP, SAS, and Halogen properties. For each generated molecule, can the authors provide information on the properties\nthat are not intended to be modified? For example, in the first subplot of Fig. S6, the generated molecules have higher QED\nvalues, but what are the LogP, SAS, or Halogen properties compared to the initial molecules?\n• If the method cannot keep most properties that are not supposed to change, the current claim of disentangled controlled\ngeneration of 3D molecules may be overstated. It would be beneficial to include sentences in the discussion detailing this\nlimitation.\n• In Fig. 2B, the UMAP analysis of the semantic embedding on QED properties shows that the semantic embeddings do not\nclearly distinguish data points with high or low QED values, even with the fine-tuned version MolDiffdAE-ft. Despite this, the\nauthors claim they can use semantic embeddings to guide the generation of molecules with higher QED properties (see Fig.\n4; Fig. 5; Fig. 6). This needs clarification. Could the authors provide more information on what they mean by\ndisentanglement in the context of the QED property?\n2. Questions about the results:\n• How do you verify that your generated structures are novel molecules and not duplicates from the training set? Additionally,\nhow do you ensure the generated molecules are valid and not just random atom substitutions?\n3. Questions about technical details\n• The atomic types are encoded as one-hot vectors in this model. Please comment on how the diffusion process is applied to\nthese vectors.\n• With such setup, it seems that you can only deal with a predefined maximum number of atoms. Does the model generalize\nwell to large molecules?\n4. Misc\n• Typo in line 6 on page 6: ‘tine-consuming’ should be ‘time-consuming’.\n• The molecule presentations in figures like Fig. 4-6 lack legends for atom types. Readers can only guess. I recommend\nindicating these types in legends, main text, or supplementary information.\n(Remarks on code availability)\nI had a brief review of the code repo. The code repo contains most of the codes of the architecture of the model and these\ncodes are well organized.\nHowever, the code repo does not have a well maintained documentation, e.g. README or doc-string in the codes. Most\nimportantly, there are no instructions on how to install this repo."}, "author_response": "Dear editor and reviewers,\nWe appreciate your efforts in reviewing our submission and your valuable suggestions. In\nthe updated manuscript, we have addressed the reviewers’ comments with updated\nexperimental results and analysis. Major updates are highlighted.\nGeneral response:\nWe would like to emphasize that unlike many of the diffusion models for 3D molecules\nlimited to random molecular generation, or a specialized task such as conditional\ngeneration or conformer generation, our model provides a generalizable framework that\ncould be adapted to many tasks that require finer controls on the generative process. We\nexperiment with benchmark tasks and also more challenging tasks that require multiple\nproperty manipulation and external datasets with limited labels. Our results show that our\nmodel could successfully control hundreds of molecular properties and manipulate up to\nthree complex molecular properties. We have verified that the molecules generated for\nthose more complicated design tasks have comparable and often superior quality with\nsimpler generative tasks (in terms of the distribution of geometric properties and energy).\nThis highlights the effectiveness of the semantic guidance, which allows for challenging\ndesign tasks without sacrificing generation quality.\nBenchmark and baseline selection\nWhile benchmark datasets like MOSES and GuacaMol have been staples for multiple de\nnovo drug design tasks, they are not suitable for our 3D-based model because they are 2D\ndatasets with no structural information, and conformer prediction methods are less\nreliable for larger molecules. As our model is specialized for the more challenging and\nhigh-dimensional task of generation of 3D molecules with both 2D and 3D control, we\nbelieve those datasets would not offer much additional insights. Instead, we utilize widely\nused benchmark datasets that contain accurate structural information, namely GEOM and\nCrossDocked, which provide sufficient information for accurate joint 2D-3D modeling of\ndrug-like molecules.\nSince the property manipulation and controlled generation in equivariant diffusion models\nis a less explored problem, we select the baselines and experiments carefully:\n1. We included a 2D VAE baseline re-trained on the same training set to demonstrate\nthe ability of autoencoding equivariant diffusion models to preserve and control the\n3D shape\n2. We also implemented classifier guidance, a traditional controlled generation\nmethod for diffusion models. Despite its effectiveness in simpler manipulation\ntasks, it fails in more complicated tasks, such as multiple properties. Our method,\nhowever, features a disentangled latent space, which enables data-efficient fine-\ntuning and manipulation of complex and potentially contradicting properties\n3. Furthermore, we did additional experiments and case studies to demonstrate the\nimportance of controlling 3D geometry, such as loosely preserving the ligand\nbinding pose\nDefinition of “disentanglement” and data efficiency\nSpecifically, we observe that our model is able to identify “disentangled”, Pareto-efficient\ndirections when manipulating multiple non-orthogonal target properties that could\npotentially conflict with each other, while traditional methods such as classifier guidance\nwould take directions that improve one objective and distract from another. This allows our\nmodel to navigate the high-dimensional, complex molecular and property space with\nmultiple constraints, and would be particularly helpful for challenging design tasks.\nOur proposed method for multiple property manipulation is highly data efficient. We use\nsmall datasets (10% of the training set for RDKit descriptors, 1000 positives for DRD2),\nresembling real world scenarios, and simple constrained quadratic optimization and\nclassifier back-propagation. Despite that, we still outperform traditional methods such as\nconditioned diffusion models and classifier guidance, which require large amounts of\nlabeled data. We believe that more effective and specialized latent space optimization\nmethods could further enhance the results.\nAdditional Experiments\nTo further prove the effectiveness of our method, we have added some additional\nexperiments:\n2.2.1: Interpretation of the latent space through direct molecular retrieval via the semantic\nembedding\n2.4.3: Enhancing the generation quality of the property manipulation tasks through retrieval-\naugmented generation (RAG). We add the results to Fig. 4 and provide a detailed case study\nin 2.4.3 and Fig. 8. These experiments demonstrate that:\n1. Our latent semantic space is well-regulated and information-rich, and allows direct\nretrieval through simple embedding comparison\n2. With RAG, even with a simple embedding fusion technique (linear interpolation), we\ncould further ensure that the manipulated molecules fall within the viable 3D\nmolecule space\nIt is also noteworthy that our method is not limited to one diffusion model framework but\ncould be potentially adapted to other models. Thus, it could benefit from the development\nof diffusion models, combining both their generative capacity and our effective semantic\ncontrol over the direction of generation. It also has the potential to be easily specialized for\nmore challenging and customized tasks (such as directly conditioning on protein pocket\nstructures), as we’ve shown with fine-tuning on external datasets and retrieval-augmented\ngeneration (RAG). We would like to leave them for further studies.\nOther important additions include:\nFig 4B: quantitative analysis of the preservation of “irrelevant” properties, further\ndemonstrating the effectiveness of disentanglement\nSupplementary 2.3: energy analysis of generated molecules. We show our model achieves\ncomparable and often superior molecular energy than the baselines. This shows our model\ndoes not compromise the generation quality despite more challenging generative tasks\nand higher model capacity.\nSupplementary 2.5: conformer generation by only updating the 3D coordinates. We show\nthe semantics control slightly improves the validity of generation\nResponse to individual reviewers’ comments:\nReviewer #1:\n1. Please include a validity analysis. It is stated that you only consider valid molecules, but\nI do not see any evaluation of how many molecules are valid for those generated.\nMoreover, the geometrical validity (e.g., bond lengths, bond angle distributions etc.) should\nbe evaluated so that readers can understand how geometrically valid molecules are. This\nis typically evaluated by other models for 3D molecule generation. For example, you could\nuse Validity3D, PoseCheck, PoseBusters and other such tools to measure this.\nThanks for the suggestion. In fact, we only report the results on valid molecules in the\ngenerated set. The success rate is the number of valid molecules that pass the similarity\nand property criteria, divided by the total number of generations.\nThe validity metric we used in the original results is whether the generated molecule\npasses RDKit’s SanitizeMol and Kekulize commands, meaning the molecules are\nconnected and the bond connection does not violate the valency constraints. We have\nclarified it in the methods section.\nAccording to your recommendations, we calculated the Jensen-Shannon divergence of the\ngeometric features (bond lengths, bond angles, dihedral angles) between the generated\nmolecules and the training set, as shown in Supplementary 2.3. We also included a quick\nenergy estimate with GFN2-xtb. Our method shows comparable and often superior\nperformance even for challenging multi-property manipulation tasks, compared to simpler\nand more stable random generation.\n2. Table 1 is not a “detailed comparison” as stated by the authors. And lacks sufficient\nconsideration of alternative methods for generation in 3D. There are many more methods\nthan a 2D autoencoder and 3D diffusion. Please include a balanced and broader\ncomparison in Table 1.\nThanks for pointing this out. We have expanded the table to include more models and\ncomparison metrics\n3. In the main text, qualitative conclusions are drawn for the fine-tuned model between\nFigure 2A left and right. This is speculative at best and to make such conclusions\nqualitative results should be included in the text. Additionally, the corresponding Figure S2\nshould include statistics so that improvements in prediction are quantifiable.\nThanks for the suggestion. We have added a more qualitative disentanglement analysis in\nSupplementary 2.2 and also the R2 values in Fig. S2.\n4. The molecules in Figure 7 are illegible and should be represented in 3D as is done\nelsewhere.\nThanks for the suggestion. We plotted the molecules with bond types and atom types to\nbetter emphasize the composition of the molecular fragments and their scoring.\nWe have proved the good preservation of 3D structures in other parts of the results. Also,\nthe placement of the atoms and bonds still follows the 3D shape of the generated\nmolecules, and the “illegible” parts are the same as the source molecule (so are naturally\nvalid) and not relevant to our analysis. We have clarified this in the figure caption.\n5. Figure 7 is also incorrectly cited as Figure S7 twice in the main text.\nThanks for pointing this out. We have corrected it.\n6. The last sentence concluding that this approach enables modelling 3D ligands without\njointly modelling the protein could be a moot point. This is because the 3D bioactive\nconformation will still need to be known in the context of binding preserving, and therefore\nthe complex structure will still need to be known. Otherwise, you risk modifying 2D\ntopology while preserving an incorrect 3D structure which is irrelevant, without guarantee\nthat the new 2D topology can form the correct 3D conformation.\n8. More important in binding preservation is the preservation of key protein-ligand\ninteractions. Only considering the docking score may slightly overlook this problem. Qualitatively, I can see examples in Figure 7 where key interactions look to be lost. This at\nleast warrants a discussion, for example, perhaps modelling polar hydrogens will improve\nthis.\nThanks for raising these valuable points. We would like to address these two questions\njointly:\nThe main goal of this analysis is to demonstrate the effectiveness, as well as the\nusefulness, of our method’s preservation of 3D topology. Since our model is trained on\nindividual molecules, there is in fact no guarantee of the preservation of interactions. Thus,\nwe show that none of the methods compared here could actively improve the protein-\nligand binding; they can only preserve it at best.\nTo fully encode, preserve, and alter the binding-related properties like the key interactions,\nthe model could be fine-tuned on the interaction complex structures with the protein\npocket as the direct condition. We would like to leave it for future investigations.\nWe have already discussed this in the Discussion section, and we also added a\nclarification in Sec. 2.4.4.\n7. Is it possible to conduct experimental evaluation of some examples to ensure that\nbinding is in fact preserved? I would expect this level of scientific rigor in Nature\nCommunications.\nWe fully agree that experimental validation is important. However, we respectfully note\nthat the primary aim of this study is to demonstrate the utility of our computational\nframework. We believe that our model, supported by rigorous evaluation and\nreproducibility, provides a solid foundation for future experimental studies and offers new\nhypotheses for further investigation. We are encouraged by the initial performance of our\napproach and are optimistic that these results will inform and guide subsequent validation\nefforts.\nWe have clarified the exploratory nature of our findings in the discussion section.\n8. The CrossDocked models are not “generally out of distribution” according to Figure\nS7(a). They overlap in chemical space and have distribution subsets that overlap with\nGEOM. It would be more robust to do a pairwise comparison of the 100\nCrossDocked test molecules to find the closest molecules in GEOM drugs using Tanimoto similarity.\nThanks for the suggestion. The point we are trying to make is that the training set GEOM,\nwhich includes isolated molecular structures, has a different distribution than the\nCrossDocked ligands. Thus, the results on the CrossDocked test set provide a solid\ndemonstration of the generalizability of the model. We have updated the description from\n“out of distribution” to “border cases”.\n9. Another potentially useful application of this approach is to preserve both the 2D and 3D\nmolecule while generating alternative conformations. Is this possible with the proposed\napproach? This could also be an effective way to measure the scope of the model.\nOur model could be used to perform conformer generation by simply clamping the atom\ntype and bond type during the inference stage, as with other diffusion models for 3D\nmolecules. The experimental results and subsequent evaluation are shown in\nSupplementary 2.5. Our method shows competitive and often superior performance on\nseveral metrics, indicating that the semantic embedding would provide some level of\ncontrol on the generation.\nWe would like to emphasize that we employed an embedding that encodes joint 2D and 3D\ninformation for more comprehensive controlled generation and property manipulation.\nConformer generation is a sub-problem that only requires 2D molecule embedding and\ndiffusion on atom coordinates only, which could be achieved with a simpler model, and is\nnot the focus of this work.\n10. Please provide specific information on the oracle used in Section 2.4.4 (versions,\nweblinks etc.).\nWe use the drd2 scorer from https://github.com/wengong-jin/iclr19-\ngraph2graph/tree/master/props. The information is included in Section 4.1.1.\n11. I don’t understand the axis of Figure S3, please add labels. I assume the x axis is 𝒛, but\nare there properties not labelled on the y axis. Why are there gaps?\nYes. The x-axis is the dimension of z. We have added the labeling to the Figure.\nThe “gaps” exist because it would be too crowded if we plotted all axis labels, so we only\nkept the important ones (key drug properties and fragment occurrences, for example). We\nhave updated the figure by annotating certain important clusters, instead of just the raw\nlabels.\n12. Please include details on training and inference time.\nWe have added the details in Supplementary 2.1.\n13. Other diffusion models have been proposed for multi-parameter optimization such as\nPILOT by Cremer et al. that haven’t been referenced or discussed.\nCremer, Julian, et al. \"PILOT: equivariant diffusion for pocket-conditioned de novo ligand\ngeneration with multi-objective guidance via importance sampling.\" Chemical Science\n15.36 (2024): 14954-14967.\nThanks for the suggestion. We have added the recommended reference in the introduction\nand discussion.\nThe README does not contain sufficiently detailed instructions for the installation and\nrunning of the code.\nWe have updated the instructions and environment files accordingly.\nReviewer #2\nWe have demonstrated the uniqueness and importance of our work in the cover letter, but\nwe would also like to further emphasize that:\n1. Unlike previous studies such as MolDiff, the focus of our work is a novel one: how to\neffectively control and steer the generation of 3D molecules with diffusion models, a\nproblem rarely investigated by other studies. Many existing works explore ways to enhance\nthe generation quality, random or conditioned, but how to control a large number of\ncompositions and properties simultaneously remains unexplored. For example, existing\nworks could generate molecules “with high QED”, “with two halogens”, “with good water\nsolubility” through conditioned diffusion models or classifier guidance, but achieving all\nthree at once remains a challenge and requires large amounts of labeled data. In our work,\nwe only use a small labelled dataset (10% of the training data) for the manipulation of up to\nthree properties, and achieve better results than traditional methods. These results\nindicate a well-regularized latent space and a strong semantic guidance of the generation.\n2. Besides small data requirements, the manipulation and RAG methods we used in this\nwork are surprisingly simple: through constrained quadratic optimization, classifier back-\npropagation and linear interpolation. Despite that, our model still achieves significantly\nsuperior results. Thus, we believe more specialized and effective methods for latent space\noptimization would be directly applied to our framework and further enhance the results.\n3. We use MolDiff as our model backbone since it is well-established and easy to compare,\nbut our framework could extend beyond one model backbone. As the capacity of diffusion\nmodels for 3D molecules grows, it is possible to combine both the higher generation\nquality and the effectiveness of our semantic guidance to further enhance the controlled\ngeneration results.\nReviewer #3:\n1. Questions related to methodology\n• If the method cannot keep most properties that are not supposed to change, the current\nclaim of disentangled controlled generation of 3D molecules may be overstated. It would\nbe beneficial to include sentences in the discussion detailing this limitation.\nThanks for pointing this out. We want to note that the “disentanglement” in the context of\ncontrolled generation refers to manipulating the disentangled features that relate to the\nmultiple and sometimes conflicting objectives, respectively. It is also worth mentioning\nthat many of the molecular properties are correlated (especially the RDKit descriptors in\nour main experiments), but it is still practically important to minimize the “side effects”\nduring the molecular manipulation.\nFor example, to increase the molecule's solubility, one manipulation approach is to\nremove the halogens, but it is not the only viable solution. Without explicit control over\nother properties, there is no guarantee which direction the modification will take. However,\nas shown in Fig. 7, when the halogen number is explicitly controlled, the model will seek\n“other modifications” that counteract the hydrophobicity of the halogen. This is an\nexample of the disentangled control: the model takes modification directions for each\nproperty that do not interfere with each other.\nAlso, as many of the properties are correlated, it is hard to select a certain group of “other\nproperties” to evaluate. Thus, we evaluate the preservation of the “other properties”\nthrough a higher-level metric of 2D and 3D Tanimoto similarity. As long as the composition\nand shape of the molecule do not change drastically, we could assume that those\nproperties are mostly preserved. We have added a new supplementary figure showing that\ncertain properties, such as the number of halogens and asphericity, would not change\ndrastically unless explicitly changed.\nWe had incorporated it into the discussion section.\n• In Fig. 2B, the UMAP analysis of the semantic embedding on QED properties shows that\nthe semantic embeddings do not clearly distinguish data points with high or low QED\nvalues, even with the fine-tuned version MolDiffdAE-ft. Despite this, the authors claim they\ncan use semantic embeddings to guide the generation of molecules with higher QED\nproperties (see Fig. 4; Fig. 5; Fig. 6). This needs clarification. Could the authors provide\nmore information on what they mean by disentanglement in the context of the QED\nproperty?\nThanks for the insightful question. In fact, the umap after fine-tuning already shows\nsignificant improvement of the separation for QED, compared to the pre-trained model.\nThe reason why it is not as visually separable as the other properties may be due to QED’s\ncomposite nature, as it is calculated from the scoring of multiple aspects of the molecule\nand may not be easily linearly separable.\nEmpirically, the claim of the successful manipulation of QED is backed by multiple\nanalyses of the real experimental results, including raw success counts and individual\ncase studies (direct QED score comparison between the manipulated and source\nmolecule). To provide a theoretical insight into the reason behind the success, our\nmanipulation method looks for a direction in the latent space that points towards higher\nQED. Thus, as long as the QED values follow a well-regulated manifold, the chances of that\nmanipulation direction ending at a higher QED value would be high.\nThe “disentanglement” in QED manipulation is similar to that of SAS and LogP in Fig. 7,\nwhere certain fragments or properties contribute to a higher QED value. In a task that\njointly increases QED and decreases SAS, for example, the fragments that increase QED\ncould either increase or decrease SAS. A disentangled approach makes sure that\nmodifications that increase QED and those that decrease SAS do not interfere with each\nother. We tried to offer a similar analysis as in Fig. 7, but the composite nature of QED\nprevents clear visualization.\n2. Questions about the results:\n• How do you verify that your generated structures are novel molecules and not duplicates\nfrom the training set? Additionally, how do you ensure the generated molecules are valid\nand not just random atom substitutions?\nThe generated molecules are filtered as follows:\nNovelty: no overlap with the training set\nValidity: (1) fully connected subgraph (2) valency constraint on each atom is not violated\nSpecifically, the validity filter is performed by the SanitizeMol and Kekulize functions from\nRDKit. These metrics ensure the reported molecules are valid and novel. We have added\nmore details in the methods section.\n3. Questions about technical details\n• The atomic types are encoded as one-hot vectors in this model. Please comment on how\nthe diffusion process is applied to these vectors.\nWe treat the one-hot vectors as multi-dimensional independent Gaussian distributions,\nfollowing practices in EDM [1] and the continuous setup of MolDiff [2]. The number of\ndimensions corresponds to the number. Empirically, a continuous diffusion model could\nstill be able to approximate the distribution of a one-hot vector rather well, as the\npredicted probability distribution on all dimensions tends to be very peaky in the later\ndenoising steps. We also show that the validity of the generated molecules is comparable\nto the baseline models.\nWe choose this diffusion process for both computational efficiency and practicality\nconcerns: a discrete representation and categorical distribution could potentially improve\nthe generation quality, as shown in [2], [3]. However, for our controlled generation tasks\nusing DDIM, such a categorical model could not be easily inverted deterministically,\nresulting in weaker controllability than a fully invertible, continuous model [4].\n[1] Hoogeboom, Emiel, et al. \"Equivariant diffusion for molecule generation in 3d.\"\nInternational conference on machine learning. PMLR, 2022\n[2] Peng, Xingang, et al. \"Moldiff: Addressing the atom-bond inconsistency problem in 3d\nmolecule diffusion generation.\" arXiv preprint arXiv:2305.07508 (2023).\n[3] Vignac, Clement, et al. \"Midi: Mixed graph and 3d denoising diffusion for molecule\ngeneration.\" Joint European Conference on Machine Learning and Knowledge Discovery in\nDatabases. Cham: Springer Nature Switzerland, 2023.\n[4] Song, Jiaming, Chenlin Meng, and Stefano Ermon. \"Denoising diffusion implicit\nmodels.\" arXiv preprint arXiv:2010.02502 (2020.\n• With such setup, it seems that you can only deal with a predefined maximum number of\natoms. Does the model generalize well to large molecules?\nThanks for pointing this out. Yes, the generation requires a pre-defined atom number.\nDuring real application, it is possible to either randomly sample or predict a viable atom\nnumber depending on the generative task.\nThe largest molecule in the GEOM-Drugs training set has 91 heavy atoms, which should\ncover most drug-like molecules. For even larger molecules (e.g. peptides, other polymers),\nthe model may need to be re-trained, but the semantic guidance methodology should be\nthe same.\n4. Misc\n• Typo in line 6 on page 6: ‘tine-consuming’ should be ‘time-consuming’.\n• The molecule presentations in figures like Fig. 4-6 lack legends for atom types. Readers\ncan only guess. I recommend indicating these types in legends, main text, or\nsupplementary information.\nThanks for the suggestions. We have corrected the typo and added an atom color coding in\nFig. 1.\nRemarks on code availability:\nWe have updated the instructions and environment files accordingly."}, {"round": 2, "reviewer_comments": {"reviewer_1": "(Remarks to the Author)\nThe authors have addressed many of the reviewer's comments, including my own, which has greatly improved the quality of\nthe work presented. I have few comments remaining.\nIn response to author comments\n> Thanks for the suggestion. In fact, we only report the results on valid molecules in the generated set. The success rate is\nthe number of valid molecules that pass the similarity and property criteria, divided by the total number of generations.\nI understand that the success rate accounts for validity by division by the total number of generations. However, it is also\nimportant to understand how much of the success rate can be attributed to valid molecule generation versus property\nimprovement. This is really a minimum requirement to report the validity (alone) of molecule generations. This should be a\ntrivial analysis to do for each of the models.\n> According to your recommendations, we calculated the Jensen-Shannon divergence of the\ngeometric features (bond lengths, bond angles, dihedral angles) between the generated\nmolecules and the training set, as shown in Supplementary 2.3.\nThank you for following my suggestion, I think the addition of this analysis is very valuable to embed it in the literature and\nbe able to compare it to previous works. Having reported these values, it seems MolDiffdAE is less accurate than MolDiff in\nreconstructing training dataset geometries, why is this? The best version with RAG results in 0.38, 0.23, 0.17, while MolDiff\nachieves 0.365, 0.155, 0.162, in bond lengths, angles, and dihedrals, respectively.\n> We plotted the molecules with bond types and atom types to better emphasize the composition of the molecular fragments\nand their scoring. We have proved the good preservation of 3D structures in other parts of the results. Also, the placement of\nthe atoms and bonds still follows the 3D shape of the generated molecules, and the “illegible” parts are the same as the\nsource molecule (so are naturally valid) and not relevant to our analysis. We have clarified this in the figure caption.\nIn that case, considering your focus on composition and following your logic, I suggest that you plot them in 2D. The current\nfigure with 3D line depictions is very messy and pointless if displaying 3D data is not needed. Moreover, I can see some\ngeometries that look strained, but I’m unsure if this is true or it is the depiction. Therefore I think it’s better to plot it in 2D,\nfor example, by using rdkit with Chem.Compute2DCoords().\n> Thanks for the suggestion. The point we are trying to make is that the training set GEOM,\nwhich includes isolated molecular structures, has a different distribution than the\nCrossDocked ligands. Thus, the results on the CrossDocked test set provide a solid\ndemonstration of the generalizability of the model. We have updated the description from\n“out of distribution” to “border cases”.\nWhether you call it “out of distribution” or “border cases” the point remains. Figure S9 does not show sufficient difference to\nprove generalizability of the model. Either test generalizability by investigating performance on external test molecules at\nbinned distances in terms of 2D / 3D distance to the nearest training dataset instance. Or simply rephrase as a case study on\nan “external test set”. Moreover, Fig. S9 is referred to twice in this paragraph when I believe it should be Fig. 9.\nAdditional comments\nRetrieval augmented generation is a nice addition to the methods. Can you please clarify if the method is used for all\nproperty manipulation tasks in Fig. 4A? It can only be seen on the top row, however, I’m not sure if it is not tested in the other\ntasks, or if it fails?\n(Remarks on code availability)\nThe README contained enough detail for installation and basic use.", "reviewer_3": "(Remarks to the Author)\nThe revised manuscript shows clear improvement, addressing most earlier concerns with new analyses and clarifications.\nThe added evaluation of irrelevant properties (Fig. 4B) and geometric validity strengthens the rigor and clarity of the work.\nOverall, the paper is significantly more convincing and well-organized than before.\nHowever, the claim of “disentangled” generation remains insufficiently explained. Although the rebuttal clarifies that\ndisentanglement refers to Pareto-efficient latent directions that minimize interference among objectives, this definition is not\nstated or fully discussed in the paper. Given that “disentangled” and “precise control” are emphasized throughout the paper,\nthe authors should explicitly define their use of this term in the Introduction and discuss it in the main text to avoid misleading\nreaders. Clarifying that the disentanglement achieved is practical rather than strict would greatly improve conceptual\naccuracy and avoid over-claiming.\nIn summary, I find the manuscript much improved and recommend it for publication after minor revision, primarily to address\nthe clarification of the disentanglement claim.\n(Remarks on code availability)"}, "author_response": "Dear editor and reviewers,\nWe appreciate your positive opinions and insightful comments. We have added additional\nmetrics to address the reviewers’ questions and qualified some claims according to the\nsuggestions. The updates are highlighted in the manuscript.\nResponse to individual reviewers’ comments:\nReviewer #1:\n1. Validity of the generated molecules\n> Thanks for the suggestion. In fact, we only report the results on valid molecules in the\ngenerated set. The success rate is the number of valid molecules that pass the similarity and\nproperty criteria, divided by the total number of generations.\nI understand that the success rate accounts for validity by division by the total number of\ngenerations. However, it is also important to understand how much of the success rate can be\nattributed to valid molecule generation versus property improvement. This is really a minimum\nrequirement to report the validity (alone) of molecule generations. This should be a trivial\nanalysis to do for each of the models.\nThanks for the insightful suggestion. We have already provided the validity comparisons for the\nRAG-based methods in Fig. 8. We have now included the raw validity counts for all experiments\nand methods in Supplementary Fig. S7. We observe no substantial differences in validity\nbetween classifier guidance and MolDiffdAE-based methods, indicating that the improvement\nin success rate is primarily driven by enhanced multi-property performance. Moreover, as\nnoted in Section 2.4.3, the performance gains achieved by RAG can largely be attributed to a\nmarked increase in validity.\n2. Geometric metrics\n> According to your recommendations, we calculated the Jensen-Shannon divergence of the\ngeometric features (bond lengths, bond angles, dihedral angles) between the generated\nmolecules and the training set, as shown in Supplementary 2.3.\nThank you for following my suggestion, I think the addition of this analysis is very valuable to\nembed it in the literature and be able to compare it to previous works. Having reported these\nvalues, it seems MolDiffdAE is less accurate than MolDiff in reconstructing training dataset\ngeometries, why is this? The best version with RAG results in 0.38, 0.23, 0.17, while MolDiff\nachieves 0.365, 0.155, 0.162, in bond lengths, angles, and dihedrals, respectively.\nWe believe the lower geometric metrics can be attributed to the following factors:\n1. Different experimental settings\nAs mentioned in 4.5.2, we have re-trained MolDiff and re-did the sampling with its\ncontinuous variant (results reported in Table S4, “DM random”), instead of the discrete\nvariant whose results (0.365, 0.155, 0.162) are reported in the original MolDiff paper.\nThis is because we employed the continuous diffusion in our work to enable\ndeterministic DDIM sampling central to our experiments, which cannot be performed\non a discrete model. Thus, a fairer comparison would be within the results reported in\nTable S4, where MolDiff random sampling showed no major advantage compared to\nRAG.\n2. Interpretation of absolute metric values.\nAll models under comparison achieve geometric metric values that are substantially\nbelow the random upper bound. Thus, the absolute differences between methods may\nnot reflect meaningful variations in geometric fidelity in practice.\n3. Trade-offs between expressiveness and unguided random generation.\nOur model employs a richer and more structured latent space, optimized specifically\nfor guided generation, including both the noise space and the semantic embedding,\nwhile MolDiff only has the former which represents a fully random latent space. As a\nresult, fully unguided random sampling (“random generation” section) is inherently\nmore challenging for our model than for MolDiff with simpler latent representations,\nand may land in infeasible regions when sampling the semantic embedding.\nHowever, under guided generation (“property manipulation” section) where the latent space is\nutilized as intended with appropriate regularization, our method consistently outperforms\nMolDiff+classifier guidance. This further highlights the enhanced informativeness of our latent\nspace, which helps constrain sampling to chemically valid regions when steering molecular\nproperties toward specified targets, especially for challenging combinations of properties.\n3. Regarding out-of-distribution samples and generalizability\n> We plotted the molecules with bond types and atom types to better emphasize the\ncomposition of the molecular fragments and their scoring. We have proved the good\npreservation of 3D structures in other parts of the results. Also, the placement of the atoms\nand bonds still follows the 3D shape of the generated molecules, and the “illegible” parts are\nthe same as the source molecule (so are naturally valid) and not relevant to our analysis. We\nhave clarified this in the figure caption.\nIn that case, considering your focus on composition and following your logic, I suggest that you\nplot them in 2D. The current figure with 3D line depictions is very messy and pointless if\ndisplaying 3D data is not needed. Moreover, I can see some geometries that look strained, but\nI’m unsure if this is true or it is the depiction. Therefore I think it’s better to plot it in 2D,\nfor example, by using rdkit with Chem.Compute2DCoords().\nWe agree with the suggestion. We have qualified the claim to “external test case” and updated\nthe figure reference accordingly.\n4. Visualization of Fig. S7\n> Thanks for the suggestion. The point we are trying to make is that the training set GEOM,\nwhich includes isolated molecular structures, has a different distribution than the\nCrossDocked ligands. Thus, the results on the CrossDocked test set provide a solid\ndemonstration of the generalizability of the model. We have updated the description from\n“out of distribution” to “border cases”.\nWhether you call it “out of distribution” or “border cases” the point remains. Figure S9 does not\nshow sufficient difference to prove generalizability of the model. Either test generalizability by\ninvestigating performance on external test molecules at binned distances in terms of 2D / 3D\ndistance to the nearest training dataset instance. Or simply rephrase as a case study on an\n“external test set”. Moreover, Fig. S9 is referred to twice in this paragraph when I believe it\nshould be Fig. 9.\nThanks for the suggestion. We have plotted the molecules in 2D in the main figure. Since the\nstructural alignments are not fully reflected in the 2D plots, we also added the 3D\nconformations in Fig. S10. We would also like to mention that the seemingly unrealistic\ngeometries are mostly caused by the orientation and 2D rendering.\n5. Scope of the RAG experiments\nRetrieval augmented generation is a nice addition to the methods. Can you please clarify if the\nmethod is used for all property manipulation tasks in Fig. 4A? It can only be seen on the top\nrow, however, I’m not sure if it is not tested in the other tasks, or if it fails?\nWe only tested the RAG in the tasks in the top row which involve more complex molecular\nproperties (QED, SAS, LogP) and more extensive comparison (with fine-tuned models and\nback-propagation-based optimization). For the remaining tasks, simpler methods already\nachieved clear improvements, so we did not include additional RAG experiments for them.\nReviewer #3:\nQualifications of “disentanglement”\nHowever, the claim of “disentangled” generation remains insufficiently explained. Although the\nrebuttal clarifies that disentanglement refers to Pareto-efficient latent directions that minimize\ninterference among objectives, this definition is not stated or fully discussed in the paper.\nGiven that “disentangled” and “precise control” are emphasized throughout the paper,\nthe authors should explicitly define their use of this term in the Introduction and discuss it in the\nmain text to avoid misleading readers. Clarifying that the disentanglement achieved is practical\nrather than strict would greatly improve conceptual accuracy and avoid over-claiming.\nThanks for the thoughtful comment. We have clarified our claims in introduction and\ndiscussions with qualifications like \"distributional regularization” and “Pareto-efficient control\ndirections”."}, {"round": 3, "reviewer_comments": {"reviewer_1": "(Remarks to the Author)\nThe manuscript authors have addressed all of my concerns in their response which I believe has greatly improved the\noverall quality of the manuscript. MolDiffdAE is a very interesting and promising approach to guided and geometrically\nconstrained chemical design and I look forward to future iterations of work in this direction.\n(Remarks on code availability)", "reviewer_3": "(Remarks to the Author)\nThe authors have addressed my earlier concerns to a reasonable extent. While I do not fully agree with all of the claims\nmade in the manuscript, the current revisions sufficiently clarify the key points, and I am willing to accept the arguments as\npresented. I have no further requests for modification.\n(Remarks on code availability)"}, "author_response": ""}]

Decision letter