MoAMa

PyTorch implementation of Motif-aware Attribute Masking for Molecular Graph Pre-training.

This version keeps the original MoAMa objective while making the training path faster, safer, and easier to transfer. The defaults now match the strongest attribute-masking configuration reported in the paper: 25% motif masking, full motif coverage, atom-type reconstruction, scaled cosine error, and a GNN decoder. See IMPROVEMENTS.md for the rationale and a suggested ablation plan.

Setup

Python 3.10 or newer is recommended.

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -r requirements.txt

Download and unpack the molecular benchmark data:

wget http://snap.stanford.edu/gnn-pretrain/data/chem_dataset.zip
unzip chem_dataset.zip

The scripts expect each dataset under dataset/<name> and its processed SMILES at dataset/<name>/processed/smiles.csv.

For ZINC2m, only smiles.csv is required. If the multi-gigabyte processed PyG tensor is absent, graphs are constructed lazily in loader workers. A bounded smoke run can be launched with --max_samples:

python pretrain.py --max_samples 10000 --epochs 1 --lazy_smiles

Pre-training

python pretrain.py \
  --dataset zinc_standard_agent \
  --output_model_file encoder

The default run uses mixed precision on CUDA, worker prefetching, AdamW, warmup plus cosine decay, residual GNN layers, and a vectorized Morgan-fingerprint similarity target. The fingerprint projector is discarded after pre-training; only the encoder is stored for transfer.

Useful switches:

# Paper-style RDKit topological fingerprint instead of Morgan/ECFP
python pretrain.py --fingerprint_type rdk

# Pure motif reconstruction ablation
python pretrain.py --reconstruction_weight 1

# Reproduce the repository's older, cheaper decoder choice
python pretrain.py --decoder mlp

--beta remains an alias for --reconstruction_weight. It now follows the paper's definition: beta * reconstruction + (1 - beta) * auxiliary.

Fine-tuning

python finetune.py \
  --dataset bace \
  --input_model_file encoder \
  --filename bace_encoder

Fine-tuning uses a lower encoder learning rate than the new prediction head, mean-plus-max pooling, a two-layer head, validation-based early stopping, and a single final test evaluation. Training history is written to log/; pass --progress when per-batch progress bars are useful in an interactive run.

The new checkpoint format records encoder architecture metadata. Legacy raw state dictionaries such as the included saved_model/encoder.pth still load; the loader expands the old chirality embedding table without altering its existing rows.

To save the best fine-tuned encoder and prediction head:

python finetune.py \
  --dataset bace \
  --input_model_file encoder \
  --output_model_file bace_finetuned