Saturn: Sample-efficient Generative Molecular Design using Memory Manipulation

TL;DR

Saturn is a novel generative molecular design method that uses memory manipulation and experience replay to improve sample efficiency, enabling potential direct optimization of high-fidelity oracles in drug discovery.

Contribution

Introducing Saturn, which applies the Mamba architecture and augmented memory techniques to enhance sample efficiency in molecular generation tasks.

Findings

Saturn outperforms 22 models on multi-parameter optimization tasks.

Memory augmentation improves sample efficiency in molecular design.

Potential to directly optimize high-fidelity oracles in drug discovery.

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.