Any-Property-Conditional Molecule Generation with Self-Criticism using Spanning Trees

## Abstract

Abstract Generating novel molecules is challenging, with most representations leading to generative models producing many invalid molecules. Spanning Tree-based Graph Generation (STGG) ( ahn2021spanning ) is a promising approach to ensure the generation of valid molecules, outperforming state-of-the-art SMILES ( weininger1988smiles ) and graph diffusion ( song2019generative ) models for unconditional generation. In the real world, we want to be able to generate molecules conditional on one or multiple desired properties rather than unconditionally. Thus, in this work, we extend STGG to multi-property-conditional generation. Our approach, STGG+ , incorporates a modern Transformer architecture, random masking of properties during training (enabling conditioning on any subset of properties and classifier-free guidance), an auxiliary property-prediction loss (allowing the model to self-criticize molecules and select the best ones), and other improvements. We show that STGG+ achieves state-of-the-art performance on in-distribution and out-of-distribution conditional generation, and reward maximization.

## 1 Introduction

Generating novel molecules is challenging, and the choice of molecular representation significantly impacts the performance of generative models. Traditional methods have mainly focused on SMILES (weininger1988smiles) 1D strings (segler2018generating; kwon2023string), and 2D graphs (jo2022score; vignac2022digress; jo2023graph). A significant issue with these representations is that a single error by the generative model can result in invalid molecules, especially as molecule size increases.

Recently, krenn2020self proposed Self-referencing Embedded Strings (SELFIES) (krenn2020self), a robust 1D string representation similar to SMILES that guarantees the generation of valid molecules through a carefully designed context-free grammar. However, recent work by gao2022sample and ghugare2023searching found that while SELFIES prevent invalid molecules, it makes exploration more difficult and reduces the performance of generative models (in terms of obtaining high-reward samples, i.e., molecules with desired properties). A significant challenge for the generative models based on SELFIES is the need to pre-define the number of tokens contained in a branch (a deviation from the main path in a 1D string) and count backward the number of tokens required to reach the beginning of the ring (starting from the end). This requires extensive planning and counting, making the problem much more challenging for the model to solve.

An alternative approach, Spanning Tree-based Graph Generation (STGG) (ahn2021spanning), has recently emerged. Unlike SELFIES, STGG is designed explicitly for generative models and works by masking invalid tokens during sampling, preventing the generation of invalid structures (e.g., atoms following atoms, branch end before branch start) and ensuring proper valency. STGG has a clear and simple set of if/else conditions used to mask out invalid tokens, which not only prevents invalid molecules but also leads to higher-quality and more diverse generated molecules (ahn2021spanning). Recent results (jang2023simple) show that STGG generally performs equally or better than state-of-the-art SMILES and graph diffusion (song2019generative; ho2020denoising; song2020score) models for unconditional molecule generation. However, its application in multi-property conditional settings has not been explored. We address this setting along with a few additional challenges, as discussed below.

Any-property-conditioning    In real-world applications, we want to be able to *generate molecules conditional on one or multiple desired properties* rather than unconditionally. Furthermore, we want to condition on *any* subset of desirable properties without retraining the model each time that we condition on a different subset of properties.

Self-criticism    Another critical issue is the synthesis time for molecules, which can take days, weeks, or even months. Thus, we cannot expect chemists to synthesize and measure the properties of all generated molecules.
Ideally, we would need a way to filter the molecules we provide to chemists. Some properties can be verified through simulations, but this can be extremely slow, and not all properties can be simulated. Another option is to rely on external property predictor models, but training, validating, and managing multiple property predictors can be troublesome. *What if the generative model could predict the properties of its own generated molecules?* This is the idea we propose here: we give the model the ability to predict properties and thus *self-criticize* its own generated molecules, allowing it to automatically filter out those with undesirable properties.

Out-of-distribution properties    Additionally, we sometimes seek to generate novel molecules with properties that have never been observed before in order to expand the range of our molecular knowledge. These *out-of-distribution properties* generally involve extreme range of values. Classifier-Free Guidance (CFG) (ho2022classifier) is a technique to improve conditioning fidelity; we found CFG useful for in-distribution properties, but problematic for some out-of-distribution conditioning values, especially for extreme values, resulting in poor generative efficiency (% of valid, unique, and novel molecules) and conditioning fidelity. Since guidance can still be beneficial to some extreme conditioning values, we propose a solution: *random guidance* with best-of-$k$ self-filtering (described further below).

In this work , we tackle any-property-conditional molecule generation with self-criticism using STGG. In doing so, we make the following contributions:

- 1.
Mixed-data Property-Conditioning: We use an MLP on standardized continuous features and embeddings on categorical features while randomly masking some properties during training, allowing conditional generation on any number of properties (0, 1, 2, or all) and the use of Classifier-Free Guidance (CFG) for improved performance.
- 2.
Improved Transformer Architecture: We use Flash-Attention, no bias terms, RMSProp, rotary embeddings, the SwiGLU activation, and better hyperparameters.
- 3.
Improved Spanning-Tree: We extend STGG to 1) allow compound structures with a new token and masking conditions, 2) prevent incomplete samples through special masking when there are too many opened branches, 3) prevent ring overflow, 4) randomize the order of the graph during training for better generalization, and 5) automatically calculate valency and adapt the token vocabulary based on the dataset.
- 4.
Auxiliary Property Prediction objective: The objective
improves conditioning fidelity and enables out-of-the-box self-filtering of molecules with incorrect properties.
- 5.
Random Guidance for Extreme Value Conditioning: Classifier-free guidance uses guidance $w>1$ to improve performance, but this can fail when conditioning on extreme values (which are needed to generate molecules with out-of-distribution properties). We propose using random guidance ($w\sim\mathcal{U}(0.5,2)$) with best-of-$k$ filtering as a solution.
- 6.
Comprehensive Performance Evaluation: We demonstrate excellent performance in terms of 1) distribution learning and diversity on unconditional generation, 2) distribution learning and conditioning fidelity on in-distribution and out-of-distribution conditional generation, and 3) diverse and high-reward samples on reward maximization.

Figure: Figure 1: Our STGG+ architecture. The molecule is tokenized and embedded. The number of started rings and embeddings of continuous and categorical properties are added, and the output is passed to the Transformer. The Transformer output is then split to produce 1) the predicted property and 2) the token predictions (masked to prevent invalid tokens). For a visual comparison between STGG and our STGG+, see Figure LABEL:fig:STGG2diag in the Appendix.
Refer to caption: extracted/5732295/STGG+.drawio.png

## 2 Background

### 2.1 1D vs 2D representations

There are various ways of representing molecules in the context of molecular generation. Some of the most popular methods are autoregressive models on 1D strings (segler2018generating; ahn2021spanning; kwon2023string) and diffusion (song2019generative; ho2020denoising; song2020score) models on 2D graphs. While both approaches have similar sample complexity, 1D strings offer a more compressed representation, requiring less space and fewer parameters and, thus, increased potential for scalability. A detailed comparison between 1D and 2D graph methods is provided in Appendix LABEL:app:1D2D.

Furthermore, recent results indicate that 1D strings are as competitive as 2D molecular graph methods for both unconditional molecule generation jang2023simple; fang2023domain and property prediction yuksel2023selformer. Both 1D and 2D representations encapsulate the same amount of information, making the choice largely a matter of preference. We advocate for 1D string representations due to their scalability and effective utilization of Transformer models, and thus, we focus on this type of representation in our work.

### 2.2 1D string representations

The most popular choice of string-based representation is SMILES (weininger1988smiles), an extremely versatile method capable of representing any molecule. However, when used in generative models, generated SMILES strings often correspond to invalid molecules. A single incorrectly placed token will lead to an invalid molecule. Graph-based diffusion methods also face a similar issue. To address this problem, recent methods like Spanning Tree-based Graph Generation (STGG) (ahn2021spanning) and SELFIES (krenn2022selfies) have been developed to prevent the generation of invalid molecules. For a detailed comparison of SMILES, SELFIES, and STGG, see Appendix LABEL:app:stggvsothers.

STGG has demonstrated performance on par with or better than state-of-the-art unconditional generative models (ahn2021spanning; jang2023simple). Conversely, SELFIES has been shown to perform worse than SMILES on property-conditional molecule generation (gao2022sample; ghugare2023searching). Therefore, in this work, we focus on STGG for property-conditional molecule generation.

### 2.3 STGG

STGG (ahn2021spanning) uses a SMILES-like vocabulary with a begin "(" and end ")" branch tokens, ring start "[bor]" and $i$-th ring end "[eor-$i$]" tokens. Contrary to SELFIES, STGG was made from the ground up for unconditional molecule generation. STGG leverages a Transformer (vaswani2017attention) architecture to sample the next tokens conditional on the tokens of the current unfinished molecule. To predict the ring end tokens, STGG uses a similarity-based output layer distinct from the linear output layer used to predict other tokens. STGG also uses an input embedding to track the number of open rings. Invalid next tokens are prevented through masking of next tokens that would lead to impossible valencies (e.g., atoms, ring-start, and branch-start when insufficient valency remains) and structurally invalid tokens (e.g., atom after atom, bond after bond, or ring-$i$ end when fewer than $i$ ring start tokens are present).

In the next section, we will show how to improve the STGG architecture, STGG, and its masking, adapt STGG for any-property conditional generation, and improve fidelity on conditioned properties through several techniques (classifier-free guidance, self-criticism, random classifier-free guidance for extreme conditioning).

## 3 Method

We tackle the problem of any-property conditional generation with self-criticism using STGG.

### 3.1 Architecture

We enhance the base architecture used originally in STGG, a regular Transformer (vaswani2017attention) directly from PyTorch main libraries. We leverage several recent improvements in Large Language Models (LLMs) to improve the architecture following GPT-3 (radford2019language), Mistral (jiang2023mistral), and Llama 2 or 3 (touvron2023llama). The improvements include:
1) RMSNorm (zhang2019root) replacing LayerNorm (ba2016layer); 2) residual-path weight initialization (radford2019language); 3) bias-free architecture (chowdhery2023palm); 4) rotary embeddings (su2024roformer) instead of relative positional embedding; 5) lower-memory and faster attention with Flash-Attention-2 (dao2022flashattention; dao2023flashattention); 5) SwiGLU activation function (hendrycks2016gaussian; shazeer2020glu); 6) changes in hyperparameters following GPT-3 (radford2019language) (i.e., AdamW (loshchilov2017decoupled; kingma2014adam) $\beta_{2}=0.95$, cosine annealing schedule (loshchilov2016sgdr), more attention heads, no dropout). These modifications aim to enhance the model’s efficiency, scalability, and overall performance.

### 3.2 Any-property conditioning

We preprocess continuous properties by applying a simple $z$-score standardization. Other transformations, such as min-max normalization and quantile transformation, were not beneficial.

To condition the model on any subset of target properties without retraining the model every time we change the properties, we need to be able to turn off the conditioning of some properties. For continuous variables, we handle missing values by introducing a binary indicator variable: if a property is missing, we set the property value to 0 and the missing indicator to 1. It is very important to include these missing indicators because we cannot assume the plausible values for the missing features (e.g., if A is 1.0, B is missing, maybe the only possible range for B is around 3 and 4, so if we leave B at 0 without missing indicator, it will not make sense). For categorical variables, we add an extra category to represent missing values. During training, the number of masked properties, $t$, is chosen uniformly between 0 0 to $T$, where $T$ is the total number of properties. This allows us to condition the model on any subset of desired properties at test time while ignoring the rest.

In the neural network, we process the standardized continuous features (continuous properties concatenated with their binary missing indicators) in a 2-layer multilayer perceptron (MLP) with Swish activation (hendrycks2016gaussian; ramachandran2017searching). Each categorical feature is then processed individually using a linear embedding. These processed outputs are added directly to the embedding of all tokens. We also experimented with injecting these embeddings through adaptive normalization (huang2017arbitrary), a method commonly used for conditioning on noise-level in diffusion models (ho2020denoising), but this approach massively increased the number of parameters without improving performance.

### 3.3 Improvements to Spanning-Tree

Starting from STGG as base, we implemented several improvements. Firstly, we extend the vocabulary to allow for the generation of molecular compounds that are composed of multiple unconnected graphs (e.g., salt is represented as [Na+].[Cl-], where [Na+] and [Cl-] are single-atom molecules connected through a ionic bond), enabling the model to solve a broader range of problems. Originally, STGG used a fixed vocabulary and a fixed set of maximum valencies that determines how many valence bonds each atom can form. Instead of requiring a predefined vocabulary, we automated the process of building a vocabulary based on the atoms found in the dataset and their maximum valency, again derived from the dataset. This data-centric approach allows us to represent complex structures, including non-molecular compounds containing metals.

We observed that STGG could occasionally generate incomplete samples by creating too many branches without closing them within the allowed maximum length, particularly when conditioning on extreme out-of-distribution properties. To address this, we modified the token masking process to ensure the model closes its branches when the number of open branches approaches the number of tokens left to reach the maximum length. This additional masking step prevents the rare but problematic situation of incomplete samples. Additionally, for massive molecules, it is possible for the model to rarely produce more rings than the maximum number of rings (100); we now mask the creation of rings when the maximum number is reached. With these additional masks, we generally maintain 100% validity, even when generating molecules with out-of-distribution properties).

### 3.4 Classifier-free guidance

To enforce better conditioning of the properties, we use classifier-free guidance, originally designed for diffusion models (ho2022classifier), and found beneficial for autoregressive language models as well (sanchez2023stay). This technique involves directing the model more toward the conditional model’s direction while pushing it away from the unconditional model’s direction by an equal amount. Figure 2 illustrates this concept. The amount of guidance typically requires hyperparameter-tuning. However, for simplicity and generality, in all analyses, we arbitrarily set the guidance parameter $w$ to $1.5$, where $w=1$ means no guidance. Practitioners can tune the guidance for further performance gain.

### 3.5 Self-criticism

To make the model more powerful, we incorporate self-criticism by adding a property-prediction loss. During training, the model is tasked with predicting both the next token in the sequence and the properties of the current unfinished molecule. During sampling, we generate molecules conditioned on the desired properties with classifier-free guidance. Then, we mask out the properties (making them fully missing) and reprocess the molecule until we reach the end-of-sequence (EOS) token. At this point, we extract the predicted property of this molecule.

This approach provides the model with a new ability: we can generate $k$ molecules for a given set of properties and then evaluate which molecule best matches its conditioned properties solely based on the model’s self-predicted properties without relying on any external property predictor. This best-out-of-$k$ strategy significantly improves the quality of its generated molecules.

The architecture with the described loss functions is illustrated in Figure 1. The methodology for generating molecules using classifier-free guidance and self-criticism is depicted in Figure 2.

### 3.6 Random guidance for extreme conditioning

Regular guidance can be problematic for some extreme (out-of-distribution) conditioning values, resulting in poor generative efficiency (% of valid, unique, and novel molecules) and conditioning fidelity. However, guidance is can still beneficial to some extreme conditioning values.

To improve generative performance on extreme conditioning values, we propose to randomly sample a guidance $w\sim\mathcal{U}(0.5,2)$ for each sample, ensuring a mix of low and high guidance. Then, our method selects the best-out-of-$k$ molecule from the molecules generated at different guidance levels, indirectly allowing the model to determine by itself which guidance is best for each sample.

Figure: Figure 2: Generation and self-prediction using STGG+. We autoregressively generate $K$ molecules conditional on desired properties using classifier-free guidance. The unconditional model predicts the properties of the $K$ molecules and the molecule assumed closest to the desired properties is returned.
Refer to caption: extracted/5732295/guidance.drawio.png

## 4 Experiments

We run four sets of experiments. First, we demonstrate that, after conditioning on in-distribution properties, our model can recover molecules close to those in the test set, achieving performance on par with state-of-the-art unconditional models. Second, we show that our model can generate molecules conditioned on properties from the test set with high fidelity on the specified properties. Third, we illustrate that our model can generate highly efficient (high % of novel, unique, and valid) molecules with high fidelity on out-of-distribution (OoD) properties. Fourth, we show that our model can produce molecules that maximize a reward function, achieving similar or better performance compared to online learning methods using offline learning. Finally, as an harder case, we show that our model can generate high fidelity molecules conditioned on out-of-distribution (OoD) properties on a small dataset of large molecules.

See Appendix LABEL:app:dataset_details for details on the datasets used. See Appendix LABEL:app:hyper for more information on the hyperparameters. See Appendix LABEL:app:prop_pred for property prediction performance metrics of the self-critic. See Appendix LABEL:app:ablation_zinc for an ablation on OOD properties for Zinc. Note that we rely on the following software: PyTorch (paszke2019pytorch), Molecular Sets (MOSES) (moses) and RDKit: Open-source cheminformatics (rdkit).

### 4.1 Unconditional generation

We train our model on QM9 (ramakrishnan2014quantum), and Zinc250K (sterling2015zinc) using the exact molecule weight, logP, and Quantitative Estimate of Druglikeness (QED) (bickerton2012quantifying) as properties. We test the similarity in the distribution of unconditionally generated molecules (masking the properties). We also test the same metrics on conditionally generated molecules conditioned on properties from the test set. We use the same train, valid, and test splits as jo2022score. We compare to the strong recent baselines reported in jang2023simple which are: EDP-GNN (niu2020permutation), GraphAF (shi2020graphaf), GraphDF (luo2021graphdf), GDSS (jo2022score), DiGress (vignac2022digress), DruM (jo2023graph), GraphARM (kong2023autoregressive), GEEL (jang2023simple), CharRNN (segler2018generating), CG-VAE (liu2018constrained), MoFlow (zang2020moflow), and STGG (ahn2021spanning).

The results are shown in Appendix LABEL:app:results_uncond. We observe that STGG+ obtains similar performance to STGG and GEEL, which are the top performing methods.

### 4.2 Conditional generation

We follow the same protocol as liu2024inverse. We train our model on HIV, BACE, and BBBP (wu2018moleculenet). We use the same train, valid, and test splits as liu2024inverse. Each dataset has a categorical property related to HIV virus replication inhibition, blood-brain barrier permeability, or human $\beta$-secretase 1 inhibition, respectively, and two continuous properties: synthetic accessibility (SAS) (ertl2009estimation) and complexity scores (SCS) (coley2018scscore). We evaluate the models using metrics on distribution and fidelity of conditioning after generating molecules conditional on properties from the test set. The condition control metrics are the Mean Absolute Error (MAE) of SAS and accuracy of the HIV categorical property (through a Random Forest (breiman2001random) predictor using the Morgan Fingerprint (morgan1965generation; gao2022sample)). The distribution metrics are Validity, atom coverage in the largest connected graph, internal diversity (average pairwise similarity of generated molecules), fragment-based similarity (degen2008art), Fréchet ChemNet Distance (FCD) (preuer2018frechet).

Notably, the Fréchet Distance is one of the most popular and meaningful distance in generative models as it correlates well with both quality and diversity (heusel2017gans); it corresponds to the Wasserstein distance on the hidden space of neural networks assuming normality. It has been used in many domains: images (heusel2017gans), audio (kilgour2018fr), videos (unterthiner2019fvd), and molecules (preuer2018frechet).

Following liu2024inverse, we compare our method to strong recent baselines: GraphGA (jensen2019graph), MARS (xie2021mars), LSTM on
SMILES with Hill Climbing (LSTM-HC) (hochreiter1997long; brown2019guacamol), and powerful graph diffusion models: DiGress (vignac2022digress), GDSS jo2022score, and MOOD (lee2023exploring), and Graph DiT (liu2024inverse). In addition, we calculate the metrics for the train and test data. This allows us to observe a cap on atom coverage in real data; therefore, coverage above the one found in the test set does not necessarily indicate better performance. Furthermore, the RandomForest predictor does not achieve perfect accuracy on the test data. While it is well correlated with better performance, any accuracy above that on the test set does not necessarily indicate better performance.

The experiments are shown in Table LABEL:tab:main2 (for the full table with more baseline approaches, see Appendix LABEL:app:fulltable). We see that STGG+ reaches near-perfect validity, coverage is consistently higher than the test set, and there is very high diversity and test-set similarity. Notably, we attain the best Fréchet ChemNet Distance; in fact, we are the only method that matches the training data’s performance, indicating that we have reached the performance cap. Regarding condition control, we achieve the best MAE on BACE and HIV while reaching the second-best on BBBP (very close to Graph DiT). Consequently, we obtain similar or better condition control than our main competitor, Graph DiT, while achieving near-optimal distribution performance (with respect to the FCD). We also obtain better performance than base STGG (with random masking and the extra symbol for compounds) on Fréchet ChemNet Distance and MAE which shows that our improvements lead to lower distance in distribution and better property conditioning.

Table: Table 1: Conditional generation of 10K molecular compounds on HIV, BBBP, and BACE.

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: In the OOD setting, how the model’s performance varies under different guidance settings, especially for extreme or non-physical property values?
soundness: 3
strengths: The paper presents molecule generation models, allowing multi-property control and self-assessment of generated molecules. It’s well-designed, with detailed experiments showing strong results across different datasets. The writing is clear and structured.
summary: The paper presents STGG+, a model for molecule generation with conditional property control. The model also has the ability of self-criticism to select optimal outputs. It achieves high validity and diversity in generated molecules, efficiently handling both typical and extreme properties.
weaknesses: The self-criticism mechanism for filtering generated molecules based on property predictions is a key feature, but there is limited evaluation of its accuracy. Detailed analysis will be necessary.

I am not an expert of this field so I will lower my confidence score.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. In SMILES notation, the same molecule can have multiple string representations. Based on my (limited) understanding of STGG, this ambiguity seems present as well. How are the molecules canonicalized?
2. In Tables 3 and A.10, do all methods reach peak performance only after generating 1 million molecules? Is the search space the same across methods?
soundness: 3
strengths: 1. To my knowledge, this manuscript is the first to thoroughly examine STGG for reward conditioning and/or optimization.
2. The work reflects a substantial effort to assess STGG+'s capabilities. Overall, the approach appears methodologically sound.
3. Molecular property optimization is an open challenge. Given its competitive performance compared to existing algorithms and its use of a (somewhat) unique molecular representation, I expect this work will attract reasonable interest.
summary: The authors build upon spanning tree-based graph generation methods to produce valid molecules with desired properties. They enhance the original network architecture by adding property embeddings and incorporate a properties predictor head for joint training. Through the use of classifier-free guidance and conditioning on these properties, the authors demonstrate that STGG+ can generate molecules conditioned on specific properties or with high reward values.
weaknesses: 1. The primary limitation of this paper is that generating 'valid' molecules does not guarantee synthesizability. Many molecules presented in the appendix would be very challenging, if not impossible, to synthesize. Meanwhile, some baseline methods may perform slightly worse on reward but produce molecules that are easier to synthesize, avoiding "reward hacking." A fairer comparison would involve evaluating the reward optimization performance of synthesizable molecules across different algorithms.
2. While novelty is not an ideal measure of a paper's value, this work is highly empirical, with limited theoretical insight. This puts more emphasis on competitive performance, yet the paper lacks adequate baseline comparisons for reward optimization. Previous work has shown that methods such as GraphGA, LSTM-HC, and Reinvent are effective at maximizing OOD reward, and these baselines should be included in Sections 4.3–4.5 (particularly Section 4.5, which currently lacks any baseline comparison). This is especially relevant as the random guidance approach for OOD generation resembles slightly enhanced random sampling with a reward proxy.
3. The molecular properties selected for optimization in this study are _very_ simple. For instance, molecular weight can be adjusted by adding or removing atoms, and logP by incorporating ionic groups (which the model does). Optimizing HOMO-LUMO gaps within the QM9 dataset is not useful, as these molecules contain only 9 atoms. These problems are generally considered solved.
4. Although the work is extensive, certain details are presented inconsistently or lack substantiation. Some claims are unsupported by data (e.g., statements like "other [configurations] were not beneficial/did not improve performance" lack any data references). Additional issues include appendix figure captions that are unclear and lack cross-references in the main text (e.g., molecules with QED > 1 in figures 8 and 14), and captions inappropriately implying low QED correlates with implausibility (e.g., figure 9). Many terms are undefined, including "synthe. MAE," "HOMO/LUMO," "SAS," as well as the precise definition of diversity used here. Additionally, error bars are missing in all tables.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: In Table 2, why not report the MAE at different percentiles, instead of only the MinMAE? It's possible that the model simply memorizes some extreme cases seen in training so as to achieve a good minimum MAE.
soundness: 2
strengths: - This paper is presented with sufficiently clear descriptions.
- The authors explored a wide range of techniques that can be applied in the under-explored context of multi-property conditional generation.
summary: This paper proposed to extend Spanning Tree-based Graph Generation (STGG) to multi-property conditional generation, namely STGG+, with improvements to successfully engineer and implement the concept. STGG+ achieves SOTA on both in-distribution and OOD conditional generation, as well as reward maximization.
weaknesses: - It seems to me that the authors invented complicated ad-hoc designs and specifically engineered to fix any issues that may arise, for example by masking the creation of rings when reaching max number (100), or alternating the use of CFG and ranking via a property-predictor. I'm afraid this hampers the overall generality of the proposed method.
- Ablation studies are missing. What's the effect of the improved Transformer architecture against the vanilla one? How does the auxiliary property prediction loss contribute to the results? The same applies to CFG w/ and w/o random guidance, the masking mechanism, the ring overflow treatment, the order randomization, and the automatic construction of vocabulary instead of a predefined one. Detailed ablations are needed to validate the authors' special designs, and provide more insight to the community.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: My questions are closely related to the weaknesses mentioned above. The authors are encouraged to address the points raised in the weaknesses section. Some questions include but are not limited to:
1. The STGG+ approach claims improvements over SMILES-based molecule generation methods. **What are some of the improvements that the SMILES representation fails to achieve or is difficult to achieve compared to the STGG+ representation?**
2. The spanning-tree representation allows both explicit and implicit hydrogen atoms. Are the generated molecules restricted to explicit hydrogen atoms or can they also have implicit hydrogen atoms?
3. How does the model/evaluation handle canonicalization of the generated sequences/molecules?
   - Does uniqueness consider canonicalization or does it only consider differences in the generated sequences?
4. For reporting the property MAEs, was an external property predictor used for the evaluation? How is MinMAE reported?
   - If an external property predictor was used, provide the details of the external predictor for MolWt, LogP, QED, and HOMO-LUMO gap.
soundness: 2
strengths: ### Originality
1. This work proposed an improved method of the STGG model for any property conditional molecule generation.
1. This work introduced classifier-free guidance and self-criticism into the transformer architecture.

### Quality
1. The proposed method is shown to improve the validity and diversity of generated molecules.
1. The method is also shown to better generate molecules with desired/conditioned properties.
1. The results are shown across multiple datasets and properties.

### Clarity
1. The STGG+ architecture is clearly explained.

### Significance
1. Any-property conditional generation is a challenging yet important task in technical applications. For instance, in drug discovery, it is important to generate molecules with desired curative properties and avoid molecules with toxic properties.
summary: This paper proposed the STGG+ method, an improved method of the Spanning Tree-based Graph Generation (STGG), for generating novel molecules with multi-property conditional generation. Architecture-wise, this work introduced

1. an improved Transfomer with Flash-Attention. RMDProp, etc.
1. an extended STGG model for more robust graph generation of molecules.

By randomly masking some properties at training time and using Classifier-Free Guidance (CFG), the model was shown to generate novel in- and out-of-distribution molecules with any property conditional generation.
weaknesses: As a computational chemist with expertise in molecules (and SMILES), I am concerned with
1. the contribution and improvement of this work, STGG+, to the original STGG model or the original SMILES-based molecule generation methods.
2. this work's representation of chemistry and molecules in terms of correctness and novelty.

**The STGG+ representation of molecules is within the capabilities of SMILES representation.**
1. The STGG+ improvements to STGG are not significant enough
   - the proposed improvements such as masking of invalid tokens, automatic calculation of valency, etc., seem similar to adding if-else conditions to improve the original STGG model.
   - In my opinion, these improvements should be learned by the model itself during training. The model itself should learn to avoid invalid tokens and keep track of valency. These are all fundamental grammar rules that the model should learn by itself.
   - If these constraints are manually added, it limits the efficiency of the generation process. For example, valency calculation can be intricate in the generation process when rings are involved.
   - Because of these manual implementations, I am not convinced that the STGG+ representation is significantly better than the SMILES representation in terms of ensuring valid molecules.
2. The proposed benefits of STGG+ (spanning-tree representation) compared to SMILES representation are not entirely true. SMILES can also achieve the claimed benefits with similar modifications. For example,
   - In SMILES, rings are represented by two identical numbers at the beginning and end of the ring. For cyclohexene, its SMILES representation can be `C1CCCCC=1` (or `C1-C-C-C-C-C=1`) and its spanning-tree representation can be `[bos]C[bor]-C-C-C-C-C=[eor1][eos]`. In the spanning-tree representation, a `[bor]` token must be paired with a `[eor#]` token before `[eos]` to form a valid molecule. Whereas in SMILES, a ring-starting number must be paired with the same number before the end of the string.
   - Automatic calculation of valency can be done in SMILES as well in the same fashion as STGG+ since the spanning-tree representation and the SMILES representation are interchangeable during the generation process.

**This work lacks clarity on the spanning-tree representation such as explicit/implicit hydrogen atoms and canonical representation.**
1. In Figure 1, line 140, the spanning-tree representation used a combination of explicit and implicit hydrogen atoms. The nitrogen atom was shown with an explicit hydrogen atom, and all the carbon atoms were shown without hydrogen atoms (implicit). However, from Appendix A.4, it seems that the vocabulary is collected for spanning tree representations with explicit hydrogen atoms.
1. The above point leads to the question of canonical representation - The same molecule can have different SMILES representations and different spanning-tree representations. In other words, different sequences of tokens can point to the same molecule. For example, `[bos]C[bor]-C-C-C-C-C=[eor1][eos]`, `[bos]C[bor]-C-CH2-C-C-C=[eor1][eos]`, and `[bos]C[bor]-CH2-C-CH2-C-C=[eor1][eos]` can all represent the same cyclohexene molecule.
1. For the reported generative efficiency (% of valid, novel, and unique molecules), was canonicalization performed/considered? If not, the reported efficiency might be overestimated. Additionally, the authors should provide some examples of the generated sequences and their corresponding molecules to clarify the canonicalization process.
2. **Suggestions:** In Section 3.3, the authors should discuss the issue with explicit/implicit hydrogen atoms and canonical representation. Try to clarify the following points:
   - Does STGG+ generate molecules with explicit hydrogen atoms only or can it also generate molecules with implicit hydrogen atoms? The spanning-tree representation should allow both.
   - How is the canonical representation handled in the training and generation processes?
   - **What is the definition of valid, novel, and unique molecules?** Are molecules considered unique if they have different sequences of tokens but represent the same molecule?

**The reported property MAE of the conditional generation needs more explanation.**
1. For the properties of the generated molecules, were they calculated with the property predictor of the STGG+ architecture or with an external property calculator such as RDKit? The external predictor should provide the ground truth for the property values and should thus be used for the evaluation.
2. The `MinMAE` reported in Table 2 needs more clarification: is it the minimum absolute error across the 2K generated molecules? What does "minimum mean" refer to?
   - If the minimum is reported, what about the mean of the absolute errors?
   - For such a large number of generated molecules, the mean absolute error is a better metric to evaluate the performance of the conditional generation. This is related to the application of the model (line 57) - validating the properties of the generated molecules in real life can be costly. Conditional generation aims to generate a small set of potential candidates. Minimum error implies that one has to test all 2K molecules (too large) to find the best candidate, while average error better represents the overall conditional generation performance.
   - The minimum absolute error might be more convincing if reported on a small population of generated molecules such as 10x molecules with multiple batches.

---

confidence: 3
contribution: 3
initial_rating: 3
presentation: 2
questions: 1. The authors improved the structure of the original Transformer, but the results do not seem to reflect the improvements, such as whether the generation time and the quality of the generated molecules have been improved.
2. For the self-criticise, the authors should discuss the trade-off between performance gains and computational cost. Including a comparison of computational time for different values of k would clarify the model's efficiency.
3. The authors should optimize the structure of the result table, as it is not clear what is being compared, e.g. modify the table head.
4. For property-conditional generation. The authors only compare the MinMSE and should add some property distributions to demonstrate that the generated molecules approximate the given conditions.
soundness: 3
strengths: This paper used a property-predictor-driven self-criticism mechanism that allows STGG+ to evaluate and select the best out of multiple generated molecules, improving fidelity to target properties.
summary: This paper proposes an enhanced version of Spanning Tree-based Graph Generation (STGG+), tailored for multi-property conditional molecule generation. Based on the STGG, STGG+ includes improvements in the Transformer architecture, a flexible conditioning mechanism for any subset of properties, and a self-criticism mechanism that filters generated molecules based on a property predictor.
weaknesses: 1. The model’s effectiveness relies heavily on the internal property predictor, which may be less reliable for out-of-distribution samples. This dependence could reduce fidelity in less representative scenarios.
2. Although the model improves conditioning performance, it’s unclear how it balances molecule diversity and property, diversity is also a crucial metric in molecular generation.

---

confidence: 2
contribution: 3
initial_rating: 5
presentation: 4
questions: Please address (to the extent possible) Q1-Q4 in the Weakness section. Specifically for Q1-Q3, please state if the alternate techniques were considered when designing the experiments/writing the paper. If so, please explain why those alternate ideas were not incorporated into the paper. If not, please discuss how these alternative techniques might be relevant to the problem considered in the paper.

Post-Rebuttal Comments
----------------------------

The authors have addressed pretty much all of my concerns. As of Nov 25, I'm keeping my original score since I'm curious to see how the reviewers with the two lowest scores (one of whom is an expert in the paper's area) respond to the author's responses to their questions.
soundness: 3
strengths: * The problem considered in this paper (molecule generation) has clear practical importance and given that there exists standard 1D string representation of molecules, using generative AI to create new molecules is definitely a  promising avenue worth pursuing.

* The paper has a pretty comprehensive experimental results and they show the efficacy of the proposed system.

* While some of the techniques used in the paper might be `standard' for say language modeling, being able to apply these techniques in a completely new application domain and show improvement is very nice.

* The ability to impose certain properties on molecules being generated and having the model be able to self-criticize seems like really nice capabilities.
summary: This paper considers the problem of generating (hopefully) new representations of molecules. Existing techniques either work with a 1D string representation of molecules or a 2D graph representation. This work considers the former representation and builds on a method called Spanning Tree-based Graph Generation (STGG) that was specifically created to use generative AI models. The present work differs from existing work along the following two high level aspects:

1. Previous work would generate molecules without any restrictions. However, this paper considers the problem of generating molecules that has to satisfy (some) subset of a given set of properties that the generated molecules must satisfy.
2. Unlike existing results the paper creates models that can _self-criticize_ by allowing the model to predict properties of the molecules it generates and uses that to prune out molecules that do not satisfy the required properties.

The paper lists conceptual improvements made in this work in the context of generating molecules but since I'm more familiar with Transformers and related literature, I will focus my review on those. Specific to the Transformer model used in this work (as opposed to the STGG work), the paper uses improvements made to the Transformer architecture (e.g. FlashAttention) over the last three years.

The paper presents a pretty comprehensive (at least to me) set of experiments and show that the proposed new system works better than existing systems on benchmarks that are used in this area.
weaknesses: Below are some questions that I was not able to answer based on what is there in the paper. (Again, as mentioned earlier, I focused in the Transformer aspects of the paper so all of the questions below are on that axis.) Also some of these questions might be asking for intuitive answers and not necessarily something that could potentially be answered with experiments-- but having these answers might be useful for the reader to understand some of the design choices made by the paper:

* (Q1) The paper uses causal Transformer models-- is there any reason a non-causal Transformer model cannot be used? E.g. non-causal Transformer models have been used on applications other than language modeling (e.g. ViT in image processing)-- and pre-Transformer language models like BERT were non-causal. In theory non-causal models are more expressive than non-causal model (just because a causal model is trivially a non-causal model as well).

* (Q2) The idea of having the model self-criticize the molecules it generates reminded me a lot of GANs. Have GANs been tried to generate molecules? If so, have they worked or is there some intuition for why they did not work?

* (Q3) Over the last ~3 years there have been a fair amount of work on `Transformer-free' models for language generation. One such line of work is based on state space model (Mamba [see https://arxiv.org/abs/2405.21060 and https://arxiv.org/abs/2312.00752] being a model that has garnered a lot of attention in the language modeling literature). Some of these ideas have been used in genomic sequencing (e.g. Hyena DNA-- https://arxiv.org/abs/2306.15794). Where these recent models considered in this work?

* (Q4) In lines 227-228, it is mentioned that the _number_ of masked properties $t$ was picked uniformly at random between $0$ and $T$. However, which of the $\binom{T}{t}$ subset of properties were actually chosen to mask?

Below are some minor comments (that are purely related to presentation):

* Lines 354-355: Instead of saying "similar" performance-- please quantify, i.e. within what percentage of existing work?

* Table 1: Is _Distance_ in the table column name the same as FCD?

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 4
questions: 1. I didn't completely follow the details of best-of-k filtering described in Section 3.6. It would be good if the authors could explain how this is exactly done. 
2. Based on observations: In Table 1, STGG+ shows an improved synth MAE, but other metrics appear comparable. In Table 2, STGG+ outperforms by design. In Table 3, the settings aren't directly comparable. Does this imply that the primary novelty lies in achieving property conditioning? If so, is the novelty somewhat limited, given that the extension feels intuitive?
3. Considering the non-comparable nature of Table 3, could the experiments be repeated under consistent settings for a fair comparison?
4. It’s currently unclear which modifications specifically drive the improvements in STGG+. Conducting ablation studies based on the points listed on page 2 (points 1-5) would provide more clarity.
5. For conditional generation, online methods like GFlowNets could serve as an additional baseline. Would it be possible to include this baseline in Table 2?

**Paper suggestions:**

a) I think, including Figure 3 (suplementary) instead of Figure 1 in the main paper would better showcase the contributions.

b) For section 3.6, I found the figure a little confusing. It may help to have a better figure to explain the same.
soundness: 3
strengths: 1. The paper is well written.
2. The authors suggest modifications that make sense and feel intuitive.
summary: This paper tackles the issue of generating valid molecules by extending the Spanning Tree-based Graph Generation (STGG) to support multi-property conditional generation. The proposed STGG+ integrates a Transformer architecture, property masking, and an auxiliary loss for model self-evaluation.
weaknesses: 1. I find the changes to be intuitive but somewhat obvious, leading me to believe that the paper lacks significant novelty.
2. Please see questions.

Meta-review: **Summary**

This work extends an unconditional molecular generation method ‘Spanning Tree-based Graph Generation’ (STGG) to enable conditional generation given some desired properties. Toward that end, the authors enhance the Transformer architecture of STGG to include (a) random masking of subsets of properties during training, (b) classifier-free guidance to generate multiple conditional samples, (c) a property-predictor component, and (d) a self-criticism mechanism to filter out from the generated molecules the ones whose predicted properties are not aligned with the desired values.  The proposed method, called STGG+, is empirically shown to achieve strong performance on reward maximization besides in-distribution and out-of-distribution (OOD) conditional generation.     


**Strengths:**

Reviewers appreciated different aspects of this work such as (a) practical relevance of the setting for molecular generation, reward conditioning, and optimization, (b) methodological soundness, (c) clear writing and presentation, (d) detailed experimental results, and (e) self-criticism as a nice capability,. 

**Weaknesses:**
Many concerns were raised by the reviewers. These included (a) validity of a generated molecule being insufficient to guarantee its synthesizability (hence the need to benchmark the methods on the performance of their respective synthesizable molecules), (b)  limited theoretical analysis or insight, (c) missing stronger baselines for OOD reward, (d) generated molecules being too simple, (e) unsubstantiated empirical claims, (f) possible over-engineering, (g) missing confidence intervals or error-bars, and (g) lack of ablation studies.  Questions were also raised about choice of some metrics and hyperparameters, conveying concerns about the generality of the proposed method as well its potential vulnerability to memorise extreme cases from the time of training. 


**Recommendation:**
During the discussion period, some concerns were addressed by the authors. However, some reviewers maintained that some major issues remained unresolved.

In particular, reviewer GMmJ asserted that the concerns about synthesizability were significant since many of the molecules generated with STGG+ seemed hard to synthesize. They also pointed to possible misalignment between reward and the ability of the different methods to generate easy-to-synthesize molecules. They also questioned whether the benefits of STGG+ were statistically significant absent the error bars.  I fully agree with these concerns. 

Similarly, Reviewer WwA3 made several critical observations. They pointed the flaws inherent in masking invalid tokens: (a) marginal performance gains did not justify the enormous overhead, and (b) STGG+ should have been able to efficiently learn to generate valid molecules intrinsically. They also pointed out to insufficiency of some experiments and reporting of experimental results: e.g., they pointed out (a) mean MAE was not shown for a statistically significant number of generated molecules), (b) values of properties such as HIV/BBBP were not based on experimental results but predicted using a random forest regressor, and (c) comparisons were not provided against  relevant evolutionary algorithms, and d) STGG+ seemed  to produce invalid generated molecules (showing charge imbalance and invalid geometries). Reviewer atbB also agreed with this assessment.  I echo these critical concerns.  

Based on the concerns raised by the reviewers, I decided to closely review the manuscript myself as well, and discovered further critical issues with the current submission.  The current state of the art in this area comprises methods such as JODO, EEGSDE, TEDMol, EDM, Modular Flows, GeoLDM, GeoBFN, and Twigs. All these methods can already achieve “any-property-conditioning”, so I’m afraid this cannot be claimed as a novelty at all. Most of these methods - with the exception of Twigs - have been around for a while, still they have not been mentioned in related work or compared against empirically. 

In fact, several of these methods are capable of generating 3D molecular conformers, beyond the capabilities of STGG+ (which is restricted to 2D generation - and this also explains in part several unresolved issues with STGG_ that many reviewers pointed out). Unless convincing experimental benefits are demonstrated against many of these strong baselines (prior published work seems to indicate that most of these baselines outperform STGG+ across metrics),  it’s hard to vote for this paper. Therefore, I recommend clear rejection.

Decision: Reject