Short-Term Graph Memory

Reference implementation for the paper Oracle-Budgeted Molecular Optimization with Short-Term Graph Memory.

The memory is a plug-in for oracle-budgeted molecular optimization. It keeps the molecules a run has already evaluated in an online GraphGPS surrogate and uses that surrogate to re-rank each round's candidate pool before the oracle is called. The oracle budget is unchanged; only the choice of which candidates spend it changes. This directory holds the memory module and the code that connects it to four generators through a single predict / update interface.

Layout

graph_memory/                        # the memory module, self-contained and importable
├── graph_mem/surrogate.py           #   SurrogateGNN: predict(smiles) / update((smiles, score) pairs)
├── graphgps/                        #   GraphGPS network and pretrained checkpoint
│   └── checkpoint/surrogate.pt
└── configs/surrogate.yaml           #   GraphGPS config, auto-located relative to the module
backbones/                           # one integration per generator, all using the same interface
├── genmol/                          #   masked diffusion (genmol_opt.py + config.yaml)
├── invirtuogen/                     #   discrete flow    (run.py)
├── reinvent/                        #   policy gradient  (run.py + hparams.yaml)
└── graph_ga/                        #   genetic search   (run.py + hparams.yaml)

The memory interface

Every backbone uses the surrogate through two calls:

from graph_mem.surrogate import SurrogateGNN
mem = SurrogateGNN(model_path="graph_memory/graphgps/checkpoint/surrogate.pt")

scores = mem.predict(candidate_smiles)      # rank a candidate pool
mem.update([(smiles, oracle_score), ...])   # fine-tune on new oracle evaluations

Self-test (run from this directory):

PYTHONPATH=graph_memory python -c "
from graph_mem.surrogate import SurrogateGNN
m = SurrogateGNN(model_path='graph_memory/graphgps/checkpoint/surrogate.pt')
print(m.predict(['CCO','c1ccccc1','CC(=O)Oc1ccccc1C(=O)O'])); m.update([('CCO',0.9)])"

How each backbone integrates the memory

All four use the GraphGPS surrogate; none use fingerprints.

Backbone	Search family	Integration file	Where memory enters
GenMol	masked diffusion	backbones/genmol/genmol_opt.py	ranks the unscored fragment pool each round; config.yaml sets surrogate.enabled: true
InVirtuoGen	discrete flow	backbones/invirtuogen/run.py	make_mem() builds SurrogateGNN, re-ranks fragment proposals
REINVENT	policy gradient	backbones/reinvent/run.py	hparams.yaml sets surrogate_type: gnn; surrogate re-ranks the sampled batch
Graph-GA	genetic	backbones/graph_ga/run.py	surrogate scores offspring before oracle evaluation

The predict / update calls are identical across the four backbones; only the surrounding generator differs. Each generator's architecture and native update rule are left unchanged. The memory reallocates the oracle budget rather than enlarging it.

Dependencies and notes

• graph_memory/ runs on its own; see requirements.txt. The GraphGPS config is located next to the module, so it works from any working directory.
• The backbones/ files are reference integrations, not standalone scripts. Each imports its own generator (GenMol, InVirtuoGen, and the mol_opt infrastructure for REINVENT and Graph-GA) and a PMO/TDC oracle. Point those imports at the model code and checkpoints from the upstream projects.
• Checkpoint paths in the backbone configs are relative to this directory (graph_memory/graphgps/checkpoint/surrogate.pt); adjust them if you run from elsewhere.
• REINVENT here uses the surrogate with hard argmax selection (hparams.yaml). For a policy-gradient generator the selected batch is also the training batch, so hard filtering can destabilize training; a soft, diversity-preserving selector is more robust. This does not affect the other three backbones.