Paper: Beam Enumeration: Probabilistic Explainability For Sample Efficient Self-conditioned Molecular Design (ICLR-2024)

Reviewer comment (Novelty):
The contribution of the paper (both conceptual and technical) is marginal.

Author response:
### **Q1: Marginal contribution and difference to Augmented Memory in Figure 1a**
The difference to Augmented Memory [1] in Figure 1a is the top box which describes the Beam Enumeration self-conditioning process (discarding molecules that do not contain extracted substructures).

Specifically, Augmented Memory is the optimization algorithm while Beam Enumeration is a general framework that extracts substructures during an optimization process for self-conditioned generation. The steps of the combined algorithm are as follows:

1. Generate a batch of molecules
2. **Beam Enumeration**: Filter batch of molecules to contain extracted substructures
3. Oracle evaluation of the filtered batch
4. Augmented Memory optimization
5. **Beam Enumeration**: If reward improves for `Patience` number of successive epochs, execute Beam Enumeration and extract new substructures

Beam Enumeration curates which generated molecules should be sent for oracle evaluation, based on the hypothesis that the most frequent (by probability) substructures are on track to yield high reward. Due to this filtering, Beam Enumeration controls what the Agent (generator) learns from, directing chemical space exploration to areas containing the extracted substructures. **We will expand the pseudo-code in the paper to also include the Augmented Memory algorithm for improved clarity.**

In the context of the combined algorithm, we would like to reinforce our contributions:

1. **Showing extracted substructures are meaningful**

In the illustrative example (Figure 2) where it is easily plausible to verify “chemical-soundness”, the results show that extracted substructures do indeed satisfy the target multi-parameter optimization objective.

2. **Early chemical space biasing is tolerable and substantially improves sample efficiency**

In the illustrative example section, we design a mechanism that decides when to execute Beam Enumeration and begin substructure self-conditioning. We believe it is not obvious that this works and would even lead to sample efficiency improvements. Consider the scenario where we are early in a generative experiment and the Agent does not yet generate molecules satisfying the target objective (for example in the early epochs of Figure 2). Beam Enumeration executes as early as the first instance in which the mean reward increases for ‘Patience’ number of successive epochs (5 was found to be optimal). This can be extremely early in a generative experiment, but we show that this immediate biasing is tolerable and substantially improves sample efficiency. To elaborate on this point, we refer to Figure C12 in the Appendix which shows the effect of Beam Enumeration filtering on the sampled batches of molecules in the illustrative experiment. At each epoch, 64 molecules are generated, and the plot shows how many are kept after filtering. At epoch 1, all molecules are kept as Beam Enumeration has not been executed. Beam enumeration executes at around epoch 5 and immediately, over half the molecules in subsequent epochs are filtered because they do not contain the extracted substructures. In the drug discovery case studies with Structure Size = 15 enforced, many epochs keep only 1/64 molecules. This is extremely biased, and we believe that showing that this is indeed tolerable and leads to substantial performance gains is a valuable contribution.

These insights lead to future research questions. For example, the empirical results show that early substructure (often scaffolds) biasing is beneficial. It is plausible that a hybrid method that aims to find high-reward substructures as fast as possible and then performing enumeration could lead to a more sample efficient method, but is outside the scope of this study.

3. **Explainability in the context of full molecules**

In the drug discovery case studies (Figure 3), we highlight the extracted substructures in the generated molecules. These substructures form central scaffolds of the full molecule and even form “uncommon” chemical groups such as the bridged cycle (top left). As molecular docking requires ligand-protein complementarity, the tolerability of these substructures provide valuable insights for structure-activity relationship studies. 

Overall, while Beam Enumeration is a heuristic, we believe the insights we draw are valuable to the ML community. The empirical performance of Beam Enumeration with Augmented Memory substantially outperforms Augmented Memory by itself which was shown in the original paper to outperform recent methods published at NeurIPS, ICLR, ICML, and AAAI (listed with their rankings in the Practical Molecular Optimization (PMO) [4] benchmark out of 30 models): **Augmented Memory** [1/30], **REINVENT** [2/30], **SynNet** (12/30) [5], **DoG-Gen** (13/30) [6], **DST** (15/30) [7], **MARS** (16/30) [8], **MIMOSA** (17/30) [9], **DoG-AE** (20/30) [6], **GFlowNet** (21/30) [10], **GA+D** (22/30) [11], **GFlowNet-AL** (27/30) [10], **JT-VAE** (28/30) [12].

Reviewer comment (Novelty):
The proposed method relies primarily on Augmented Memory [1], and the methodology outside of Augmented Memory (*Beam Enumeration* ~ *Self-conditioned Generation* paragraphs in Section 3) is short and seems like an additional trick.

Author response:
Beam Enumeration demonstrably improves performance over state-of-the-art methods like Augmented Memory [1] and REINVENT[2, 3]. The increased sample efficiency and extractability of interpretable substructures represent empirically measurable progress in areas of importance to the molecular design community. The tangible improvements over strong baselines highlight the usefulness beyond an incremental "trick".

Reviewer comment (Novelty):
What is the difference of the generation scheme in Figure 1a with that of Augmented Memory?

Author response:
### **Q1: Marginal contribution and difference to Augmented Memory in Figure 1a**
The difference to Augmented Memory [1] in Figure 1a is the top box which describes the Beam Enumeration self-conditioning process (discarding molecules that do not contain extracted substructures).

Specifically, Augmented Memory is the optimization algorithm while Beam Enumeration is a general framework that extracts substructures during an optimization process for self-conditioned generation. The steps of the combined algorithm are as follows:

1. Generate a batch of molecules
2. **Beam Enumeration**: Filter batch of molecules to contain extracted substructures
3. Oracle evaluation of the filtered batch
4. Augmented Memory optimization
5. **Beam Enumeration**: If reward improves for `Patience` number of successive epochs, execute Beam Enumeration and extract new substructures

Beam Enumeration curates which generated molecules should be sent for oracle evaluation, based on the hypothesis that the most frequent (by probability) substructures are on track to yield high reward. Due to this filtering, Beam Enumeration controls what the Agent (generator) learns from, directing chemical space exploration to areas containing the extracted substructures. **We will expand the pseudo-code in the paper to also include the Augmented Memory algorithm for improved clarity.**

Reviewer comment (Presentation):
The *Autoregressive Language-based Molecular Generative Model* paragraph in *Proposed Method* section should be moved to *Related Work* section as that part is not proposed by this paper.

Author response:
### **Q2: Autoregressive Generative Model section moved to Related Work**

We agree that autoregressive generative models is related work but we included it in the `Proposed Method` section to introduce token sampling and how goal-directed generation can be framed as a reinforcement learning (RL) problem. One crucial finding we want to convey in our work is that *during* an optimization process, Beam Enumeration extracts meaningful substructures. Thus, we included this section to introduce the optimization algorithm.

Reviewer comment (Novelty):
The assertion that this is the "first method to concurrently address explainability and sample efficiency" is overstated. Guo et al. [1] have put forth a concept strikingly akin to the one proposed here, where they too derive interpretable substructures and then generate novel molecules. Likewise, Fu et al. [2] have purportedly tackled both interpretability and enhanced sample efficiency.

Author response:
**Response Part 1/2**

We thank the reviewer for the critical feedback and pointing to existing works that provide notions of interpretable substructures and sample efficiency. We intended to convey that Beam Enumeration is the first method that shows sample efficiency and explainability can be made synergistic. We will modify our statements in the main text to reflect this and to better convey that Beam Enumeration shows that sample efficiency and explainability can be made **synergistic**. In our response, we first highlight what we mean by synergistic. Next, we discuss specific differences between Beam Enumeration and the works by Guo et al. [4] and Fu et al. [5].

**Beam Enumeration displays synergistic behaviour between sample efficiency and explainability**

The founding hypothesis of Beam Enumeration is that during an optimization trajectory, it must become increasingly likely to generate favorable molecules and this is why extracted substructures are meaningful. In our experiments, goal-directed generation is achieved by explicit multi-parameter optimization (MPO) via reinforcement learning. We designed Beam Enumeration to execute when the reward improves for `Patience` number of successive epochs (in which we find 5 is optimal). The key information Beam Enumeration provides is which substructures are on track to receiving a high reward from the oracle. We want to emphasize that the substructures satisfy the oracle because it means that if the oracle is replaced by a higher-fidelity oracle (more computationally costly simulation or even experiments) such as molecular dynamics simulations [6], the extracted substructures would become more meaningful. Given a more sample efficient method, the reward would improve faster, and often more drastically. This is something observed in the Augmented Memory [3] paper but was not pointed out by the authors of Augmented Memory. Specifically, Figure 2a in the paper shows that the Augmented Memory optimization trajectory rapidly improves with a steep gradient compared to other methods. Since Beam Enumeration extracts substructures when the reward improves, it is expected that a more sample efficient method can actually lead to more informative substructures extracted earlier. This is an explanation of why Beam Enumeration improves Augmented Memory to a greater degree than REINVENT; because Augmented Memory is more sample efficient. **By nature of being more sample efficient, the substructures become more meaningful, which in turn, further improves sample efficiency.** Augmented Memory is the state-of-the-art sample efficient model in the Practical Molecular Optimization (PMO) [7] benchmark, and the addition of Beam Enumeration substantially improves its sample efficiency.

**Response Part 2/3**

**Difference to Data-Efficient Graph Grammar (DEG)**

Guo et al. [4] introduce the DEG method which, given a dataset, constructs a graph vocabulary that can reconstruct the dataset. With the constructed vocabulary (where building the molecule is referred to as applying production rules), reinforcement learning can be applied to tune the probabilities of production rules and optimize for an objective. The first distinction of Beam Enumeration is that DEG is inherently constrained (by design) to generate molecules with the learned Graph Grammar. Moreover, we note that their case studies focus only on polymers, which are constructed from monomer building blocks, which fits their method well. **The relevant question in goal-directed generation is, what if no molecules in the training dataset can satisfy the objective function?** In such a case, the distribution of the generative model must be shifted. As an example of this, we refer to the MK2 task in Table 2 in the Beam Enumeration main text, where the pre-trained model is insufficient in generating high-reward molecules. Beam Enumeration with both Augmented Memory and REINVENT shifts the distribution to accomplish the task, as shown by the Yield and Oracle Burden metrics improving. With a constrained vocabulary, DEG may not be as performant when the task requires generating molecules de novo (no known examples of active molecules), as the Graph Grammar construction deeply depends on the initial dataset. This is in contrast to pre-trained general models like Augmented Memory and REINVENT, which construct molecules token-by-token. Beam Enumeration extracts substructures also token-by-token, which enables access to diverse substructures. Concrete examples of this are shown in Figures D13-15 in the Appendix, where some substructures are particularly “uncommon”, which may carry even more valuable insights.  **Moreover, sample efficiency was not discussed in the DEG work**, and the optimization objectives were diversity and Retro* Score, which can be generally more lenient oracles. Given a relevant drug discovery objective to optimize physics-based oracles, it is unclear how well DEG performs with a constrained vocabulary. We emphasize that we are not claiming DEG cannot perform this task, as another popular generative model, JT-VAE [8] also generates a notion of “vocabulary” via junction tree subgraphs. The notable result is that pre-training on a large dataset can construct a (relatively) larger vocabulary to access a larger chemical space (but still unable to generalize to any arbitrary molecule, as stated in the paper). However, the question still remains how sample efficient these models with constrained vocabularies would be, as the JT-VAE for example, is ranked 28/30 in the Practical Molecular Optimization (PMO) [7] benchmark while Augmented Memory is the state-of-the-art and REINVENT is ranked 2. Beam Enumeration substantially improves the sample efficiency of both these methods.

**Response Part 3/3**

**Difference to Differentiable Scaffolding Tree (DST)**

Fu et al. [5] introduces the DST method which can yield interpretability and explicitly discusses sample efficiency. The interpretability in DST comes from the local gradients which can provide insights into which chemical group attachment positively or negatively affects the oracle evaluation. Therefore, extracting interpretability involves looking at sets of local gradients to draw insights, which is less accessible. By contrast, Beam Enumeration extracts entire scaffolds and substructures, which then self-condition the generative process going forward such that **every single molecule that is kept must contain one of these substructures.** This effectively equates to a guided exploration of chemical space where generated molecules are in some sense “enumerated” based on central substructures and scaffolds. Provided an informative oracle, this can yield structure-to-activity relationship insights. We explicitly discuss this in the main text in Figure 3, when we present the extracted substructures in the drug discovery cases and show the tolerability of large chemical groups in the binding sites. 

While the DST paper also discusses sample efficiency, there are notable differences. Firstly, DST is ranked 15/30 in the Practical Molecular Optimization (PMO) benchmark [7] (Augmented Memory is 1st and Beam Enumeration substantially improves it). Secondly, the sample efficiency was evaluated on surrogate predictors and relatively more permissive oracles (for example logP, where simply adding more carbon atoms would increase it). Importantly, DST “saves” on oracle evaluations by training a GNN via supervised learning to mimic oracle evaluations. **This has a crucial effect in generative design, as the GNN should only be expected to be accurate in the chemical space it is trained on, which inherently restricts the chemical space.** Very recent work [9] has explored how a surrogate model can be used to mimic oracle calls in a generative paradigm and the solution is that the surrogate model does not actively predict the label, but rather, chooses which molecule to send for oracle evaluation (essentially acting as a curator). For this reason, it is unclear how well DST would perform in an optimization task involving physics-based oracles (as is used in real-world case studies) where the oracle is not permissive and surrogate model inaccuracies can quickly degrade performance. Beam Enumeration with Augmented Memory **explicitly** optimizes for physics-based oracles, and the empirical evidence shows that it is performant in these tasks.

We thank the reviewer for raising their score and for providing further feedback. 

* We removed all mentions of “first method to jointly address explainability and sample efficiency”. In the latest uploaded PDF, we added a passage to acknowledge DEG and DST in the `Explainability for Molecules` section in `Related Work`. The following passage was added:

**“In the context of generative models, previous works have explicitly addressed explainability [1, 2] and jointly with sample efficiency [2].**

Reviewer comment (Novelty):
Beam search is a well-established technique within the NLP community. I would recommend revisiting the nomenclature for the proposed method.

Author response:
Beam Enumeration is distinct from Beam Search [10, 11]. Here, we describe the similarities and differences point by point: 

• **Objective:** Beam Search approximates finding the highest probability sequences by maintaining a "width" of the top k candidates at each step. The number of sequences decoded via Beam Search is exactly k. By contrast, Beam Enumeration exhaustively enumerates the top k tokens at each step to extract substructures, such that the total number of sequences is $k^{beam-steps}$.

• **Search Space**: Beam Search only keeps the top k candidates globally, pruning away unlikely candidates. Beam Enumeration keeps the top k local candidates per sequence to do exhaustive enumeration of sub-sequences.

• **Termination**: Beam Search terminates when a sequence ends or a maximum length is reached. Beam enumeration does pre-defined steps of sub-sequence enumeration.

• **Use of results**: Beam Search returns the top candidates to use directly. Beam Enumeration uses the enumerated sub-sequences to extract molecular substructures for self-conditioned generation.

• **Sequence Length**: Beam search produces complete sequences. Beam enumeration produces sub-sequences of predefined length.

In summary, Beam Enumeration does local, exhaustive enumeration of sub-sequences for a different end purpose compared to Beam Search - to extract meaningful substructures rather than find a few optimal sequences. The exhaustive sub-sequence enumeration enables probabilistic extraction of knowledge from the model's built-in likelihoods. In the main text, we emphasize the difference in the `Proposed Method` section in the following line:

“We note the closest work to ours is the application of Beam Search for molecular design [12]  to find the highest probability trajectories. Our work differs as the objective is not to find a small set of the most probable sequences. Rather, we exhaustively enumerate the highest probability sub-sequences to extract molecular substructures for self-conditioned generation.”

We proposed the “Beam Enumeration” nomenclature to encapsulate the algorithm behavior which is to enumerate the beams, but it may lead to confusion as pointed out by the reviewer. To more clearly differentiate Beam Enumeration and Beam Search, we have added the following sentence at the end of the `Beam Enumeration` introduction section in `Proposed Method: Beam Enumeration`:

**“We further detail the differences between Beam Enumeration and Beam Search in Appendix G”.**

Correspondingly, the newly added Appendix G contrasts the behavior, similar to our previous response.

Reviewer comment (Evaluation):
Would it be possible to exhibit the SA score, as per reference [1], for the molecules generated in the context of drug discovery (Figure 3)?

Author response:
The SA score (calculated according to [13]) is now annotated in Figure 3 (**we have updated the manuscript PDF**) for all generated molecules and the reference. The SA score of the generated molecules are similar to the reference, which is expected, as QED [14] was used in the optimization objective. The SA score is based on fragment contributions from molecules in PubChem, which contains many “drug-like” molecules and QED explicitly optimizes for Lipinski-compliance. We would like to note that while the SA scores of the generated molecules are similar to the reference, we do not make any claims on synthesizability as the SA score assesses molecular complexity rather than explicitly synthesizability. Moreover, if SA score were the objective to be optimized, it could be included in the MPO objective.

Reviewer comment (Experiments):
Have the authors considered employing a more advanced language model, like the transformer, as an alternative to LSTM?

Author response:
Transformers [15] are ubiquitous in NLP and possess excellent generative capability, particularly highlighted by recent LLM applications. However, our decision to use a RNN-based model is for the following reasons:

**1. Short token-sequence lengths**

For many drug design tasks, the focus is on small molecules (and is our focus in the paper) which are generally defined as being < 500 Da. Consequently, the token-sequence lengths are typically < 60 which is in stark contrast to NLP transformer settings where entire passages and pages can be passed as context. These short sequences mean that the limitations of RNNs are not pronounced. The subsequent points build on this observation.

**2. RNNs satisfy base metrics of molecular generative models**

We refer to base metrics as **Validity** (generates valid molecules), **Uniqueness** (generates unique molecules), and **Novelty** (generates molecules outside the training data). Early benchmarks such as GuacaMol [16] and MOSES [17] have shown that RNNs operating on SMILES achieve high values on these metrics and are therefore a solved problem.

**3. RNNs learn the training data distribution**

Pre-training the generative model requires learning the underlying training data distribution. Specifically, pre-training involves learning the conditional token sampling probabilities that reproduce the training data. [18] shows that RNNs can learn these distributions well, and even for irregular distributions.

**4. Transformers are not necessarily better than LSTM-RNNs in reinforcement learning (RL)**

We first acknowledge existing works using transformers for SMILES-based molecular design [19, 20, 21]. It was observed in [21] that transformers were unstable during RL optimization and more prone to mode collapse (generating the same molecules repeatedly).

Next, we draw insights from the RL literature. [22] showed that transformers can struggle in certain optimization settings, often not outperforming random generation. [23] made similar observations and proposed an architectural modification to enable effective RL. Recently in [24] (NeurIPS 2023), the authors derive insights into *when* transformers may be less performant than LSTM when using RL. Their results showed that in short-term memory tasks, transformers had worse sample efficiency compared to LSTMs. Therefore, based on the literature and empirical observations, we decided to use a LSTM-RNN model. We end our response by acknowledging that very recent work [25] (NeurIPS 2023) has shown that transformers can also be performant in molecular design tasks with RL and is an area we will explore in the future.

Reviewer comment (Evaluation):
Is there a direct comparison between the presented methods and those in PMO using the original metrics?

Author response:
In the Augmented Memory [3] work, their PMO [7] results show that it is the new state-of-the-art. However, the tasks in the PMO benchmark do not involve optimizing physics-based oracles, which are ubiquitous in real-world molecular design and are the oracles that **demand** sample efficiency since they are computationally expensive. For example, **almost all works that have achieved experimental validation using goal-directed (explicit optimization of a MPO objective) generative design used physics-based oracles [26-31].** This is also the case for the current most advanced clinical candidate (phase 2 clinical trials) designed by a generative model [32].

For this reason, we focus on such tasks in our evaluation of Beam Enumeration. We make an effort to show broad applicability by performing 3 drug discovery tasks, each involving different protein receptors, which naturally require different ligand-protein complementarities. The results show that adding Beam Enumeration on top of Augmented Memory substantially improves the sample efficiency. Our objective (together with the code release) is to promote practical molecular design, especially in settings that demand sample efficiency.


We believe this work provides an important advancement in integrating explainability into goal-directed molecular generation, while enhancing sample efficiency. The ability to extract substructures that improve performance could provide crucial insights for medicinal chemists and reduce time spent in the discovery process. The reviewer's valuable feedback has helped refine our ideas and improve the presentation. We hope our responses satisfactorily address the reviewer's concerns. We thank the reviewer again for their time.

Please let us know if you have any further questions!
\
\
Sincerely,
\
\
The Authors

Reviewer comment (Evaluation):
There are additional results in the appendix regarding the impact of the proposed Beam Enumeration scheme on reducing diversity. While the main text briefly mentions such diversity reduction is expected, the paper would benefit from having (at least some) further discussion (in the main text) on this effect and its extent.

Author response:
We thank the reviewer for their positive assessment of our work and for valuable feedback. We will respond to each of the reviewer’s questions individually.

### **Q1: Beam Enumeration’s effect on diversity**

We have expanded the `Quantitative Analysis: Sample Efficiency` section to include information about the diversity. Specifically, we have added the following sentence:

“Beam Enumeration decreases the diversity of the generated set (as measured by IntDiv1 [1]), but finds considerably more unique scaffolds above the 0.8 reward threshold (up to 19x) (Appendix D). With many unique scaffolds built around (often) central substructures, the generated set could conceivably provide insights into structure-activity relationships.”

We previously omitted this information due to space constraints but we agree with the reviewer (and thank the reviewer for pointing this out) that this information is important, as it imposes some trade-off when using Beam Enumeration.

Reviewer comment (Experiments):
There should be further investigation of the impact of k (i.e., the number of top substructures) on the performance of Beam Enumeration. The authors hypothesize that a low top k would be optimal, but wouldn't the use of small k potentially lead to molecules with repetitive substructures and undesirably limit the molecular diversity?

Author response:
We believe there may have been a typo in the reviewer’s comment as `top k` controls the number of top probable tokens to enumerate at each time step, while the `Pool Size` controls the number of top substructures to track. We kindly ask the reviewer to let us know if we have misinterpreted the question.
For completeness in our response, we provide our rationale for the choice of both these hyperparameters.

The `top k` hyperparameter controls the top number of most probable tokens to enumerate at every time step during Beam Enumeration. **The grounding hypothesis of Beam Enumeration is that the most probable substructures (which are extracted from sub-sequences of tokens) are on track to yield high reward.** We chose to use a `top k` of 2 because 1 would essentially equate to greedy decoding (just taking the top probability at each time step which would only yield 1 sequence). By contrast, enumerating the top 2 tokens at each time step enables generating an exhaustive set of token sub-sequences. Our decision to not use a higher `top k` comes from three observations:

1. A higher “k” increases the likelihood that the extracted substructures are no longer the most probable and our hypothesis is that the most probable substructures are those that are on track to yield high reward. This observation is elaborated in subsequent points.

2. There is a GPU memory overhead of exhaustively enumerating sub-sequences. By choosing `top k` = 2, we can enumerate for longer beam steps. Longer beam steps means that the token sub-sequences are longer, which would equate to extracting larger substructures. Our results show that larger substructures are vital to further enhancing sample efficiency. Finally, we note that while we mention there is a GPU memory overhead, the added wall time is negligible compared to the oracle calls (we show all wall times in Appendix D.4.)

3. If we choose a higher `top k`, it is likely that the extracted substructures are more diverse. However, our results suggest that this is disadvantageous for improving sample efficiency. In Appendix A, we show exhaustive hyperparameter investigations (over 2000 experiments) for Beam Enumeration. One hyperparameter in particular is the `Token Sampling Method` which takes on two possible values: “topk” or “sample”. “Topk” enumerates the top k highest probability tokens at each beam step. By contrast, “sample” samples k tokens from the conditional probability distribution such that on average, the highest probability token is chosen, but there is a chance that another token is chosen. In our experiments, we find that “sample” sampling can sometimes be more performant than “top k” sampling, **but exhibits notably higher variance**. This higher variance means “sample” sampling is less stable which is in agreement with our hypothesis that the highest probability substructures are those that carry the most meaning, as “sample” sampling can extract less probable substructures. For these reasons, we decide to use a “top k” of 2.

Building on the observation that the most probable substructures are the most informative, we decided to use a `Pool Size` of 4 (keep the top 4 substructures) as the same hyperparameter investigations in Appendix A show that it is, on average, the most performant. We tried `Pool Sizes` of [3, 4, 5].

Taking all the results together, Beam Enumeration can substantially improve sample efficiency with a small trade-off in diversity.

Reviewer comment (Experiments):
Similarly, the manuscript could benefit from having further discussion on the optimal choice of the "Structure Minimum Size" - currently set at 15 in the experiments - and its impact on the overall performance. It is mentioned that "larger substructures ... improves performance" but what is the best minimum size to use? Increasing it beyond some number would certainly degrade the performance of the molecules and also limit the diversity of the generated molecules significantly.

Author response:
### **Q3: Impact of `Structure Minimum Size` on performance**

In our initial experiments, we explored substructure extraction either by `Scaffold` (extract Bemis-Murcko scaffold) or `Structure` (extract any valid substructure). When `Structure Minimum Size` was not enforced, `Structure` extraction often led to small, functional groups as the most frequent substructures. Concrete examples of these substructures are shown in Figure C12 in the Appendix. We find that this is the source of the performance discrepancy between `Structure` and `Scaffold`, where the latter improves sample efficiency more (Table 1 in the main text). This observation reaffirmed our hypothesis that extracted substructures are indeed meaningful, otherwise, it is highly unlikely that the `Scaffold` substructures, which are larger, would improve sample efficiency. Based on this, we hypothesized that enforcing `Structure Minimum Size` could put `Structure` on par with `Scaffold`. We verify this statement in Tables 1 and 2 in the main text. Our response so far is to share our decision making when designing Beam Enumeration. 

In this second part of the response, we explicitly answer the reviewer’s question. Since, the evidence so far suggested that larger substructures improves sample efficiency, we intentionally included `Structure Minimum Size` (no minimum size, 10, and 15) as part of the exhaustive hyperparameter screening (Figures B5-10 and Tables 6-11 in the Appendix). **Our findings are that higher `Structure Minimum Size` always improved sample efficiency.** This is shown in the 3 drug discovery case studies (Table 2). Examples of these substructures are shown in Figures D13-15 in the Appendix. **From this set of results, we present the first conclusion: Beam Enumeration with higher (up to 15) `Structure Minimum Size` enforcement improves sample efficiency with a trade-off in diversity.** As it is the most performant across the illustrative and the 3 drug discovery experiments (with replicates), it is robust as the default hyperparameter value and is designated as such.

Next, we answer the reviewer’s question about further increasing `Structure Minimum Size`. There are two things we considered: 

1. **GPU memory overhead.** To obtain larger substructures in our current framework, the beam steps would need to be increased so that the token sub-sequences are long enough to generate larger substructures. This comes with a computational bottleneck as the number of sub-sequences is given by top k^beam_steps which grows exponentially. As such, there is a practical limit (assuming modest computational resources) to brute force enumeration. One solution is to prune some sub-sequences by some criteria such as if the next token is below some probability. This would “free” some memory to continue enumerating other sub-sequences. However, this leads to our second consideration.


2. **Signal-Noise ratio.** Beam Enumeration extracts the most frequent substructures from the exhaustive set of token sub-sequences. If we begin pruning sub-sequences to enable some longer sub-sequences, the exhaustive set is going to be smaller. This means that the frequency ranking of extracted substructures may be more noisy, such that the performance exhibits higher variance.

These are open questions that we will explore and have deferred for future work as our objective was to introduce the Beam Enumeration framework and show that it leads to improved sample efficiency and provides explainability out-of-the-box.

Reviewer comment (Experiments):
The multi-property optimization (MPO) aspects will need to be discussed in further details. Currently, the MPO appears to completely rely on the baseline algorithm (e.g., Augmented Memory or REINVENT) to be extended with the Beamn Enumeration capability. But since different substructures might be associated with different properties, incorporation of Beam Enumeration for MPO would have to consider the impact of self-conditioned molecular generation when used with a specific baseline algorithm and a specific number of properties for optimization. For example, one may expect that when optimizing a large number of properties, then using a very small k may be detrimental as the selected substructures may be associated with only some of the properties of interest.

Author response:
### **Q4: MPO optimization and substructure bias being detrimental to sample efficiency**

In both the illustrative experiment (optimize topological polar surface area, molecular weight, and number of rings >= 2) and drug discovery case studies (QED score, molecular weight, docking), the objective functions are MPO objectives. The results in the paper therefore show that Beam Enumeration with both Augmented Memory [2] and REINVENT [3, 4] can handle MPO settings. However, we supplement our response with other challenging scenarios that are often encountered in real-world drug discovery efforts and how the Beam Enumeration bias does not have to be detrimental to sample efficiency.

**1. Simultaneous optimization of many properties**

The reviewer brings up a good point about substructures potentially being associated with different properties. In the experiments in the paper, all individual properties to be optimized are combined into a single scalar value during reinforcement learning optimization and this can handle the MPO nature of the problem. However, it is possible to put more weight on certain properties (for example, docking may be more important than molecular weight). By doing so, the substructures would naturally be more biased towards the docking optimization since the weight updates of the Agent depend on this oracle feedback (in which docking now has more onus). This is something that has been done in REINVENT [4] and has shown success. 


**2. Orthogonal MPO objectives**

We further discuss another scenario where properties to optimize may be orthogonal (they impose a trade-off). Optimizing property 1 may impose a trade-off on property 2 (pareto-front problem). This can be encountered in real-world drug discovery problems such as designing for affinity towards 1 protein while discouraging affinity for another protein. A solution here is to apply curriculum learning [5]. One could decompose the MPO objective into sequential objectives, for instance, to first optimize for affinity to protein 1 and then optimizing to discourage affinity for protein 2 (negative design). Beam Enumeration would then extract substructures in line with each individual objective, removing the problem of “conflicting signals”. However, the substructure themselves would always need to be interpreted relative to each other, when extracting insights.

Overall, there are scenarios where the MPO nature of the problem may impose difficulties in optimization (not just for Beam Enumeration, but generally). However, there are procedures and established techniques that can enable Beam Enumeration to be applied without conflicting substructures.

Reviewer comment (Novelty):
Finally, while the paper focuses on language models for molecular generation and uses Beam Enumeration for self-conditioning the generative process by filtering the most probable substructures, it would be great if the authors could discuss how similar ideas may be used to extend other types of popular generative molecular design models (e.g., latent space models based on VAE, JT-VAE). Especially, models like the JT-VAE optimize molecules by sampling & assembling substructures (in the form of junction trees), which - at least at a conceptual level - *may* be related to the core "self-conditioning" idea in Beam Enumeration, to prioritize certain substructures. While this may be beyond the scope of the current work, having at least some high-level discussion might benefit the readers.

Author response:
### **Q5: Applying Beam Enumeration beyond language-based formulations**

In the first part of our response, we highlight some notable differences in the specific formulations of VAE [6] and JT-VAE [7] models that may prevent **direct out-of-the-box** application of Beam Enumeration. In the second part of the response, we discuss how **concepts and findings** from Beam Enumeration can be used to formulate a combined algorithm with JT-VAE. 

VAEs and the JT-VAE learn a latent space encoding of the training data and decode to generate new molecules. One difference is that to achieve goal-directed generation (explicitly satisfy a target objective), an optimization algorithm needs to be applied. In the JT-VAE paper, bayesian optimization is applied to optimize for logP. While showing the capability for optimization, logP is generally not representative of real-world drug design tasks since simply adding carbon atoms would increase the logP. Based on the results, it is unclear how well the optimization process is for multi-objective optimization with physics-based descriptors. Beam Enumeration is built on the observation that it must become increasingly likely to generate favourable molecules if the optimization process is successful. This leads to the second difference, particularly for the JT-VAE. During training, the **JT-VAE learns substructures as a function of the training data** and is not guaranteed that a junction tree exists for an arbitrary molecule (as stated in the paper). This inherently constrains the generalization of the model, especially in the often case where desirable molecules are out-of-distribution to the training data. By contrast, **Beam Enumeration learns substructures as a function of the oracle**, which means that as long as the oracle is optimized over time, extracted substructures are meaningful for that oracle (which should be designed to be informative and correlated with the desired objective). Out-of-distribution generation is handled innately since molecules are generated token-by-token, which allows very different molecules to be generated to the training data. A concrete example of this is in the MK2 case study in Table 2. Very few molecules are generated in the base algorithms which indicate the pre-trained model is not particularly well-suited to satisfy the MK2 docking objective. **The addition of Beam Enumeration to Augmented Memory shifts the distribution and finds substructures that satisfy this oracle.**

We next discuss how **concepts and findings** from Beam Enumeration can be applied to JT-VAE. JT-VAE shows that junction trees can be used to construct local neighbourhoods. Given a favourable seed molecule (a molecule that satisfies the target objective reasonably well), using JT-VAE to propose a local neighbourhood expansion and then scoring these molecules would set-up the possibility to iteratively perform substructure biasing. Specifically, taking the set of all molecules in the local expansion, the most frequent substructures could be used for the next epoch of generation. This could potentially build up larger substructures (sub-graphs) over time.

We end our response by providing an example where Beam Enumeration could be applied out-of-the-box to graph generative models. GraphINVENT [8] is an autoregressive graph model where molecular generation proceeds by “graph action probabilities” which denote whether to add a node or edge. [9] couples reinforcement learning to GraphINVENT, which sets up the potential to directly apply Beam Enumeration by enumerating the most probable graph actions and extracting the substructures. However, based on the PMO [10] benchmark results and recent work on state-of-the-art graph generative models [10, 11], SMILES-based LSTM models are often more sample efficient and better capture the training data distribution. For this reason, we focus on Augmented Memory and REINVENT which are the most sample efficient models in the PMO benchmark.

**We have added a section (Appendix F) discussing the application of Beam Enumeration beyond language-based models.**

Reviewer comment (Experiments):
It would be great if authors could incorporate more baseline methods.

Author response:
### **Q1: More baseline methods**

In the Augmented Memory [1] work, their PMO [2] results show that it is the new state-of-the-art in sample efficiency (REINVENT [3, 4] being ranked 2nd). For this reason, we focused on applying Beam Enumeration to Augmented Memory and REINVENT to demonstrate that their sample efficiency can be further improved. Our goal in the present work is to push the limits of sample efficiency and demonstrate synergy with explainability by proposing the general Beam Enumeration framework.

However, Beam Enumeration is not limited to Augmented Memory and REINVENT. We have added a section (Appendix F) which discusses how Beam Enumeration could be extended to graph-based autoregressive models and other sub-graph methods. In the future, we will explore extensions of Beam Enumeration to other classes of generative models, and other language-based architectures.

Reviewer comment (Writing):
Some minor issue, e.g., repetitive statements, like "To address this, Gao et al. (Gao et al. (2022) proposed the Practical Molecular Optimization (PMO) (Gao et al. (2022) benchmark." Too many "Gao et al".

Author response:
### **Q2: Repetitive statements**

We have removed the redundant “Gao et al.”.

Reviewer comment (Experiments):
Do you use the same setup with PMO? If not, it seems unfair to compare with PMO's best baseline REINVENT.

Author response:
### **Q3: Same set-up as PMO**

In the Augmented Memory [1] work, the authors benchmarked their model on the PMO [2] benchmark, following the same set-up (10,000 oracle calls and across 23 benchmarking tasks). In the Beam Enumeration manuscript, we do not benchmark performance on the PMO benchmark. Our reasoning is because none of the tasks in the PMO benchmark involve optimizing physics-based oracles (such as molecular docking). These are the oracles that **demand** sample efficiency since they are computationally expensive. Importantly, **almost all works that have reported experimental validation** of a generated molecule used physics-based oracles [5-10]. This is also the case for the current most advanced clinical candidate (phase 2 clinical trials) designed by a generative model [11]. For this reason, we focused on optimization tasks involving physics-based oracles. We make an effort to show broad applicability by performing 3 drug discovery tasks, each involving different protein receptors. The comparisons between Augmented Memory [1], REINVENT [3, 4], and the addition of Beam Enumeration to both, are fair, as all experiments enforced an oracle budget of 5,000 calls.

Reviewer comment (Reproducibility):
Is it possible to incorporate MD-based simulation as oracle? How long does it take for each oracle call? Can you provide more details?

Author response:
### **Q4: MD-based simulation as oracle**

MD-based oracles are much more computationally expensive compared to molecular docking. For example, free energy perturbation [12] which is much more accurate than docking in predicting binding affinity, can take 6 GPU hours per molecule. A generated batch of molecules could be parallelized across multiple GPUs. **However, we want to be prudent in making definitive claims without empirical evidence and thus, our answer is that Beam Enumeration with Augmented Memory [1] is not sample efficient enough.** 

This is also one of the motivations of introducing the Oracle Burden metric. With these higher-fidelity oracles, we may not need to generate a large number of ideas to triage, as is the standard practice currently; optimizing for molecular docking and then selecting a subset for MD simulations. With higher-fidelity oracles, generating even 10-100 favourable molecules may be sufficient. Beam Enumeration is a step towards achieving this aspiration.

We believe Beam Enumeration is a valuable general framework to extract explainability and improve sample efficiency in current state-of-the-art (based on sample efficiency) generative models. We want to thank the reviewer again for their time and feedback.
\
\
Please let us know if you have any further questions!
\
\
Sincerely,
\
\
The Authors