SEGO

SEGO (Sample-Efficient Generative Optimization) is a language model-based framework for de novo small molecule design. It combines autoregressive molecular generation with Bayesian optimization to maximize oracle scores with minimal oracle calls.

Installation

Requirements: Python 3.10, CUDA-enabled GPU (CPU works but is much slower), Linux.

# 1. Install Conda, then clone the repository
git clone <repo-url> && cd sego

# 2. Create and activate the environment
conda env create -f segobabel_env_nobuilds.yaml
conda activate segobabel

Common installation issues

• GLIBCXX_3.4.29 not found:

  conda uninstall openbabel
  conda install gcc_linux-64 gxx_linux-64
  conda install -c conda-forge openbabel

• causal-conv1d / mamba-ssm build errors: see Issue #1.

Running SEGO

SEGO is launched via saturn.py with a JSON configuration file:

python saturn.py config.json

Configuration

The JSON config controls all aspects of the run. Key parameters under reinforcement_learning:

Parameter	Description
sego	true to enable SEGO's BO-guided inner loop
inner_loop_mode	"oracle_anchored" (oracle-supervised reset), "continuous" (full memory), or "fresh" (clean slate each round)
inner_loop_rounds	Number of rounds where the generator is optimized against the surrogate
sampled_library_size	Molecules sampled from the agent for selection by the GP
inner_loop_diversity_filter	Whether to apply diversity filtering inside the inner loop

Default configs are in configs/: saturn.json for a plain augmented-memory run and gollum.yaml for standalone Gollum BO. Additional example configs and pretrained checkpoint models are in experimental_reproduction/.

Acknowledgements

Saturn code adapted from:

• Saturn (preprint)](https://arxiv.org/abs/2405.17066) — git fee0179

Gollum code adapted from:

• Gollum (preprint) (https://arxiv.org/abs/2504.06265) - https://github.com/schwallergroup/gollum

Mamba architecture code adapted from:

• Official Mamba repository
• PTM-Mamba
• Mamba CPU