HGR: Higher-Order Molecular Grammars for Generative and Foundation Models in Chemistry

This repository contains the official implementation of HGR.

Yiming Huang, Yujie Zeng, Vijay Prakash Dwivedi, Simone Foti, Jianmin Wang, Jure Leskovec and Tolga Birdal

🌐 Project page

Figure 1. HGR lifts a molecular graph into a combinatorial complex, induces reusable production rules, serializes them as an HGR string, and reconstructs the molecule by reverse derivation.

HGR turns hierarchical molecular topology into a sequence of grammar production rules. Rings and motifs become explicit higher-order structures in a combinatorial complex, while the resulting rule sequence remains compatible with standard sequence models. The same representation is used for molecular generation and transferable property prediction.

Code

Component	Purpose	Main entry point
HGR grammar	Construct MIG or RSG grammars and reconstruct molecules	scripts/construct_grammar.py
HGR-VAE	Learn a continuous latent space over grammar-rule sequences	scripts/gvae_train.py
HGR-LDF	Generate molecules by diffusion in the HGR-VAE latent space	scripts/main_diff.py
HGR-FM	Pretrain and transfer a grammar-based molecular encoder	scripts/fm_pretrain.py, scripts/fm_finetune.py
RingDiv	Evaluate generation on ring-enriched molecular data	ringdiv/

Installation

Create a Python environment, then install HGR and the RingDiv evaluation package in editable mode:

python -m pip install -e .
python -m pip install -e ringdiv/

Data and checkpoints

📦 Datasets and pretrained checkpoints: coming soon.

Runtime data stays outside Git history. To keep datasets, checkpoints and outputs under the repository root, use:

export ASSET_ROOT="$PWD"
mkdir -p datasets checkpoints results logs

You can point ASSET_ROOT elsewhere if the assets live on another disk.

Running generation experiments

Paper configurations are grouped by dataset under configs/. Commands below use repository-relative config paths.

# Build a grammar and train HGR-VAE
python scripts/construct_grammar.py --config configs/qm9/gvae_rsg.yaml
python scripts/gvae_train.py --config configs/qm9/gvae_rsg.yaml

# Sample from a trained HGR-VAE
python scripts/gvae_sample.py \
  --config configs/qm9/gvae_rsg.yaml \
  --ckpt /path/to/gvae_checkpoint.pth

# Train and sample from HGR-LDF
python scripts/main_diff.py \
  --config configs/zinc250k/diff_rsg.yaml \
  --mode train
python scripts/main_diff.py \
  --config configs/zinc250k/diff_rsg.yaml \
  --mode sample \
  --ckpt /path/to/diffusion_checkpoint.pt

# Evaluate generated molecules
python scripts/eval_gen_rel.py \
  --config configs/qm9/gvae_rsg.yaml \
  --smi_path /path/to/generated.smi

Final generation configs are available for QM9, ZINC250k, RingDiv300k, MOSES and GuacaMol. Sweep definitions remain in the internal research repository.

Running foundation-model experiments

The foundation model uses one pretraining config and one downstream config:

python scripts/fm_pretrain.py --config configs/pretrain.yaml
python scripts/fm_finetune.py --config configs/finetune.yaml

The downstream config performs full fine-tuning by default. To evaluate with frozen-encoder probing, add --ft-type freeze.

Configuration and assets

Every --config argument accepts either an absolute path or a path beginning with configs/, resolved from the repository root. Checkpoint arguments accept absolute paths or paths relative to $CKPT_ROOT, which defaults to $ASSET_ROOT/checkpoints.

configs/env_config.yaml is not part of the public release. Set machine-specific paths through ASSET_ROOT and optional service settings such as WANDB_ENTITY through environment variables.