Saturn: Sample-efficient Generative Molecular Design using Memory Manipulation

Authors: Jeff Guo, Philippe Schwaller
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: generative design, drug discovery, sample efficiency, language models, reinforcement learning, scaling

Abstract:
Generative molecular design for drug discovery has very recently achieved a wave
of experimental validation, with language-based backbones being the most common
architectures employed. The most important factor for downstream success is
whether an *in silico* oracle is well correlated with the desired end-point. To this end,
current methods use cheaper proxy oracles with higher throughput before evaluating
the most promising subset with high-fidelity oracles. The ability to directly optimize
high-fidelity oracles would greatly enhance generative design and be expected to
improve hit rates. However, current models are not efficient enough to consider such
a prospect, exemplifying the sample efficiency problem. In this work, we introduce
**Saturn**, which leverages the Augmented Memory algorithm and demonstrates the
first application of the Mamba architecture for generative molecular design. We
elucidate *how* experience replay with data augmentation improves sample efficiency
and *how* Mamba synergistically exploits this mechanism. Saturn outperforms 22
models on multi-parameter optimization tasks relevant to drug discovery and may
possess sufficient sample efficiency to consider the prospect of directly optimizing
high-fidelity oracles.

Decision: reject
Reviewer scores: 6, 6, 5, 5, 3

Reviews and discussion:
<thread>
creation date:
1730506760878

thread_type:
official_review

<thread_object>
creator: reviewer_f1gd

<content>
summary:
This paper propose a molecular generation framework for drug discovery by combining language model and Mamba neural network. The author experimentally show that a better generative performance can be achieved by incorporating Mamba network into the language model.

strengths:
The empirical evaluation is comprehensive. The author also provide open-source implementation, which enhance the credibility of the paper.

weaknesses:
1. The presentation of the paper should be significantly improved. The description of the proposed framework is hard to understand. The author should at least provide a description of how the molecular generation problem is formulated as an optimization problem using clear math notation in Section 3. Many terminologies, such as "local sampling", "augmentation round", lacks clear description. The author should also include a clear description of Mamba network as most of the reader may not be familiar with it. 
2. It is unclear what the novelty of this paper is. The novelty of this paper seems to be limited to the application of Mamba network to molecular generation. 
3. Many model choices/claims of the proposed framework lack justification. For example, the author assume "oracle evaluations are near deterministic" without any justification. The author also need to discuss why Mamba can be potentially better compared to transformer/RNN. 
4. The experiment part of this paper seems to be simply Methods X Metrics. It is unclear what the author want to show other than the empirical superiority of Mamba compared to RNN/transformer. The author claim more than 5000 experiments by taking different random seeds into account, which is quite misleading. The author also need to compare the computational consumption of Mamba compared to RNN/transformer. 
5. The math notation is messy. For example, in Eq. (1): what is $x$? what is $\theta$?

questions:
1. Could the author provide a more principled description of "augmented memory"? The explanation in Section 3 is hard to understand. 
2. The Saturn perform really bad compared to GEAM and the author improve Saturn by finetune the model with dissimilar data. Do the authors do the same for GEAM?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730589458206

thread_type:
official_review

<thread_object>
creator: reviewer_u4yg

<content>
summary:
The paper proposes a novel framework that introduces the Mamba architecture and enhances sample efficiency through augmented memory and experience replay mechanisms, demonstrating superior performance on multi-parameter optimization tasks relevant to drug discovery compared to the baseline models.

strengths:
This paper presents the first application of the Mamba architecture in generative molecular design. It uses augmented memory and experience replay to significantly enhance sample efficiency, enabling faster optimization of high-fidelity oracles. The authors conducted sufficient experiments to demonstrate the model's generative capabilities.

weaknesses:
1. While Saturn excels in identifying high-reward molecules, there is a trade-off in diversity, which may affect its capability for exploratory molecular design.
2. Although this paper used Mamba for the first time in the molecular generation task, the innovation of the model is limited.

questions:
1. While the authors have demonstrated the efficiency of the Mamba through the comparison of different models under the Oracle Burden (OB) condition, it would be great to evaluate and report the actual generation time for each model. Could the authors provide any insights or results regarding the runtime performance of the different models?
2. The authors should clarify why they specifically chose to maintain only the top 100 molecules in replay buffer. Additionally, does the replay buff lead to the potential issue of storing very similar molecules in the buffer, which could subsequently reduce the overall diversity of the generated samples?
3. In part 3, the authors only compare the metrics related to Hit Ratio. Is there a problem that the generated molecules may be similar? Some metrics related to the diversity of generated molecules should be added for comparison.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730677415602

thread_type:
official_review

<thread_object>
creator: reviewer_xhfy

<content>
summary:
Motivated by the need to optimize expensive oracle functions, the authors propose Saturn, a framework for sample-efficient molecular optimization. Saturn consists of the Mamba architecture for generating SMILES, and RL with Augmented Memory for optimization. Extensive experiments show promising results on optimizing docking score compared to baselines in a sample-efficient manner.

strengths:
* The motivation of the work is very clear and convincing. Inaccurate in silico oracles is one of the major limitations of generative models in practice, and so being able to utilize highly accurate but costly oracles in generative models would be very useful.
* Clear and well-written paper, including methods and results.
* Very extensive results. Comparison with ~20 state-of-the-art molecular generative models on the docking task, where Saturn appears to be significantly superior to all except GEAM.
* Meticulous ablations of each model component and architecture hyperparameters. The appendix is also very in-depth.

weaknesses:
* Very limited technical novelty. Saturn seems to be the application of the previously developed Augmented Memory approach, except swapping the backbone for the Mamba architecture. Thus, while the paper contains strong empirical results, it is probably not the right fit for ICLR.
* The paper is motivated by the need to optimize the scores from high fidelity oracles, but no experiments are conducted for this setting. The results on docking are promising, but it is not clear that the superiority of Saturn will transfer to the more complex case of optimizing high fidelity oracles. For a paper that is mostly empirical, it seems important to have results on the motivating example instead of a toy example.
* The most competitive baseline, GEAM, seems to reach about similar performance to Saturn for most tasks. While the authors argue that GEAM achieves more diversity at the cost of less sample efficiency, it appears the authors did not explore tweaking the hyperparameters of GEAM. Is it possible some hyperparameters of GEAM might encourage less diversity and increase sample efficiency?

questions:
* The authors state that they add a Tanimoto dissimilarity objective to Saturn to make Saturn-Tanimoto, but I don’t see a similar GEAM-Tanimoto. Wouldn’t it be possible to add the same optimization to GEAM? If so, would it be more fair to compare Saturn-Tanimoto with GEAM-Tanimoto?
* What causes the improved sample efficiency of Saturn? It’s not clear that the Mamba architecture would be particularly sample efficient, so it seems to be the SMILES data augmentation method. However, I don’t see this explained in depth anywhere in the main text, particularly how the augmentation is actually done.
* It is not clear to me in the methods which contributions are yours and which come from previous works. Could this be clarified?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730699345516

thread_type:
official_review

<thread_object>
creator: reviewer_feff

<content>
summary:
- This paper tackles the problem of sample efficiency in chemistry generative models.
- It proposes to use experience replay and data augmentation to train a Mamba model and show its sample efficiency compared to other state-of-the-art molecular generative models.

strengths:
- The manuscript is well written and was easy to follow. The appendix is extensive and provides enough details on each component of the method. 
- The experiments are well thought and extensively demonstrate most of the laid out claims.

weaknesses:
- The authors have claimed that promised candidates from docking are subjected to MD simulation which is a higher fidelity oracle and direct optimization for higher fidelity oracle may be possible with sample efficient generative models (Saturn). However, this paper remains short of demostrating the utility of Saturn on MD simulations and leave it as future work. It is not clear if the sample efficiency gains observed in case of a comparatively lower-fidelity oracle (docking) would translate to high-fidelity oracle (MD). For the sake of completeness and justification of laid out claim, I would suggest that the authors provide some preliminary results on application of Saturn for MD.

questions:
- On the effect of diversity filter in augmented memory: Does setting the reward to 0 and removing all instance of over-represented scaffolds in the buffer, limits the generator from learning from potentially useful scaffold, if the removed scaffold had high reward before being penalized for redundancy?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730902607618

thread_type:
official_review

<thread_object>
creator: reviewer_smbd

<content>
summary:
The authors present Saturn, an RL-approach using a Mamba backbone for generative molecular design. Design choices surrounding implementation are thoroughly explored and motivated towards increasing sample efficiency to enable the use of higher fidelity (and more computationally expensive) objective functions.

strengths:
The authors focus on a weakness of the current trend in generative chemistry to couple an optimization algorithm with an insilico oracle function, suggesting two avenues for improvement: (i) increase the fidelity of the fast oracle, or (ii) increase the sample efficiency of the generative process to enable more costly oracles at reasonable throughput. In this paper they focus on (ii) which is a valuable contribution to the field.

The authors previously described the RL paradigm and here focus on substituting the language model backbone within the constraints that it is sufficiently simple to fine-tune on the fly. Given the recent success of Mamba as a linear time sequence modelling architecture this choice seems well motivated and suitable.

The reinforcement learning paradigm including SMILES augmentation, and experience replay buffer update with a genetic algorithm is intricately described along with its effect on the diversity of generated outputs.

weaknesses:
Major:
- The main weakness of the method seems to be the diversity reduction exhibited upon moving to Mamba from an RNN backbone. Therefore the majority of my main questions are surrounding this observation.
  - The authors describe a logical diversity filter on the augmented memory buffer that uses bemis murcko as a filter. I am surprised that this addition does not seem to mitigate the reduction in diversity. Could the authors comment on this, and possibly show the murcko scaffold diversity along with their results in Table 2.
  - Following the argument at Line 292 that Mamba generates repeat SMILES because its loss is lower during pre-training, this argument would suggest overfitting. If the authors use the Mamba checkpoint with a loss matched to the RNN, is the performance restored to the RNN levels of diversity? At what sample efficiency cost?

Minor:
- The introduction should mention the use of QSAR predictive oracles since these are very common in addition to docking.
- The paper is motivated by an intent to use molecular dynamics as an oracle but does not achieve this goal, could the authors add a comment on what it would take to achieve this in the discussion section.
- The authors discuss the trade off of sample efficiency and diversity, this is similar to the common fidelity vs diversity tradeoff among image generative models. The authors might consider a 2D summary plot showing the tradeoff of these while varying batch size and n augmentation rounds since its a lot of information to digest in a tabular form.
- I am not clear on the logic presented in the section “squeezing the likelihood of Augmented SMILES” the authors might consider clarifying.
- Figure 2d compares the trajectory of generation throughout augmentation rounds. It is unclear that the Mamba path is more linear from these plots in my opinion. (i) The UMAP should be fit to the same set in order to compare and (ii) the RNN should also have 10 rounds.

questions:
- Is there a possibility that the molecules generated by GraphGA impair the quality of the replay buffer molecules as this is an imperfect generative approach itself, would it be better to use a higher fidelity method to breed the members of the buffer.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731667968769

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Manuscript Changes

comment:
### **Manuscript Changes**

All changes are highlighted in the updated manuscript version.

* Additional sentence introducing QSAR models in the Introduction **following Reviewer SmbD’s feedback**

* Fig. C3 in the Appendix showing a 2D heatmap of sample efficiency and diversity trade-off for Mamba with batch size 16 (Saturn uses this). **This follows Reviewer SmbD’s feedback**

* Contributions list previously mentioned performing > 5,000 experiments. **Following Reviewer f1gd’s feedback**, we changed this to “> 500 experiments, all across 10 seeds)”

* Appendix C.2 description of how SMILES augmentation is performed **following Reviewer XhFy’s feedback**

* Appendix F.2 wall times for reproducing GEAM **following Reviewer u4yG’s feedback**. Saturn’s wall times were previously already reported in Appendix F.1. Saturn’s generation time is faster. 

* Page 4, clarity on what “augmentation round” means **following Reviewer f1gd’s feedback**.

* Page 4, added sentence about fixing the docking oracle seed which to support our statement about deterministic oracles **following Reviewer f1gd’s feedback**

* Page 5, clarity on what “local sampling” means **following Reviewer f1gd’s feedback**

### **New Results (attached zip file)**

* Showing direct high-fidelity oracle optimization. Density Functional Theory (DFT) oracle optimization - plots and discussion attached as a supplementary file

* Further analysis between Saturn and GEAM results showing that GEAM generates molecules with lower ligand efficiency and are larger, which can inflate docking scores. This is due to GEAM not optimizing QED to the same extent as Saturn.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732125290973

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
We are eager to discuss with reviewers

comment:
As half the discussion period has passed, we wanted to acknowledge that our responses are long but we wanted to be thorough and would be grateful if reviewers can consider the points we raised. **We only want to reiterate the following points:**

### **1. GEAM's pre-training requires a labelled dataset**
All ~250k molecules in ZINC 250k were docked (across all 5 protein targets). ***Therefore, GEAM pre-training requires an up-front cost of 1.25 million docking oracle calls (250k molecules, 5 protein targets).*** GEAM then trains a model separately for each protein target. This is in stark contrast to Saturn's unsupervised pre-training, which requires no up-front oracle calls and the pre-trained model is applied as is, to all protein targets. The shared sentiment amongst reviewers is that GEAM's performance is similar to Saturn. ***We just want to note that we are evaluating Saturn at an extreme disadvantage***. Despite this, Saturn still optimizes the objective to a ***much greater*** degree (**see Strict Hit Ratio results**).

### **2. GEAM’s generation scheme**
***Always*** starts from benzene and adds fragment attachments to construct the full molecule. This means it cannot generate molecules without benzene. However, due to their use of GraphGA directly as their genetic algorithm, final generated molecules *may* not have benzene due to crossover and mutation operations in the GA. Nonetheless, the constraint of always starting from benzene is a practical limitation. 

### **3. DFT optimization**
Given the above points, we hope to have conveyed to the reviewers that Saturn considerably outperforms GEAM. We go further to show preliminary results on direct DFT optimization (**attached supplementary file**). We believe most models are not close to this prospect of direct high-fidelity oracle optimization.

\
\
We are eager to engage in discussion with the reviewers to improve our work and incorporate feedback.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This submission introduces Saturn for generative molecular design based on the language model backbone Mamba. Saturn integrates augmented memory in RL for more efficient generation with strong empirical performance in drug discovery. 

The reviewers agreed that the performed experiments, including the ones during rebuttal, demonstrated the sample efficiency and generation performance by Saturn. However, most of the reviewers consider that the technical contributions are limited given both Mamba and augmented memory mechanisms are not new. The authors in the rebuttal explained that the novelty lies in elucidating "the deeper architectural and optimization insights in the context of molecular design".

The authors may consider the review comments for possible future submission, better positioning the research work, presenting the technical contributions in more self-explaining and rigorous ways, and providing stronger evidence of how the design of Saturn and obtained insights may help further improve molecular design.