Graph Energy Matching (GEM)

Official repository for Graph Energy Matching: Transport-Aligned Energy-Based Modeling for Graph Generation.

This v0.1 code release focuses on unconditional molecular graph sampling from the MOSES dataset. It covers:

• transport-loss training on MOSES
• transport + contrastive fine-tuning
• unconditional MOSES FCD/metrics
• unconditional MOSES sampling animations

For the continuous-data Energy Matching codebase, including a 2D toy Jupyter notebook, see EnergyMatching and its toy example.

Setup

conda env create -f environment.yaml

Data (MOSES)

MOSES CSVs are downloaded automatically on first use and cached under data/moses/moses_pyg/.

Checkpoints

The exact MOSES checkpoint used for the paper/demo will be released soon.

Train Transport Loss

Start from scratch and use transport loss only:

python src/gem/train_gem_ebm_fm.py \
  --config-name gem_ebm_fm_moses_ver3.yaml \
  train.max_iters=390000 \
  train.batch_size=128 \
  fm.batch_size=128 \
  train.save_every=30000 \
  train.lambda_fm=1.0 \
  train.lambda_cl=0.0 \
  train.cl_steps=0 \
  train.resume=null \
  train.init_ckpt=null \
  viz.enabled=false \
  viz.run_during_training=false \
  ++sample.evaluate=false

Tune Mixing Proposal Sampler

After transport-loss training, calibrate the mixing proposal sampler so it can traverse the learned energy landscape. The resulting parameters can be used for contrastive fine-tuning.

python src/gem/cali_params.py \
  --config-name cali_params_moses.yaml \
  train.resume=/path/to/transport/checkpoint.pt \
  train.init_ckpt=null \
  cali.search.n_trials=100 \
  cali.N_noise=32 \
  cali.N_data=0 \
  cali.chain_warmup.steps=225 \
  cali.steps=200

For a new transport checkpoint, replace the calibrated sampler values below: train.dl_beta, train.dl_lambda_X, and train.dl_lambda_E.

Fine-Tune Transport + Contrastive

Start from a transport-loss checkpoint:

python src/gem/train_gem_ebm_fm.py \
  --config-name gem_ebm_fm_moses_ver3.yaml \
  train.init_ckpt=/path/to/transport/checkpoint.pt \
  train.resume=null \
  dataset.filter=true \
  train.max_iters=500 \
  train.lr=1e-5 \
  train.batch_size=128 \
  fm.batch_size=128 \
  train.save_every=100 \
  train.lambda_fm=1.0 \
  train.lambda_cl=0.1 \
  train.cl_warmup_iters=100 \
  train.cl_loss_threshold=0.2 \
  train.cl_steps=200 \
  train.proposal=dlangevin_vec \
  train.dl_beta=8.76576 \
  train.dl_lambda_X=0.326323 \
  train.dl_lambda_E=2.36417 \
  train.gamma_train=0.5 \
  train.chain_warmup.enabled=true \
  train.chain_warmup.steps=225 \
  train.chain_warmup.proposal=simple_ver2 \
  train.chain_warmup.simple_n_edits=1 \
  viz.enabled=false \
  viz.run_during_training=false \
  ++sample.evaluate=false

Use train.init_ckpt to start fine-tuning from transport weights. To continue an interrupted fine-tune run, use train.resume=... train.init_ckpt=null.

Evaluate FCD And Metrics

This writes metrics_over_time.json, metrics_over_time.csv, and generated SMILES tables:

python src/gem/metrics_over_time.py \
  --config-name gem_metrics_over_time_moses.yaml \
  metrics_run.checkpoint=/path/to/model.pt \
  metrics_run.total_samples=25000 \
  metrics_run.batch_size=512 \
  metrics_run.fcd_reference_size=25000 \
  metrics_run.steps='[0,50,100,200,350]'

MOSES unconditional animation

python src/gem/animate_energy.py \
  --config-name gem_ebm_animation_moses_public.yaml \
  viz.checkpoint=/path/to/model.pt

Citation

If you use this code, please cite the Graph Energy Matching paper. The MOSES dataset is from Polykovskiy et al. (2020).

Parts of the graph transformer, MOSES data processing, molecular metrics, and discrete flow-matching utilities are adapted from DeFoG (Qin et al., 2025).

@article{balcerak2026graphenergymatching,
  title={Graph Energy Matching: Transport-Aligned Energy-Based Modeling for Graph Generation},
  author={Balcerak, Michal and Shit, Suprosanna and Prabhakar, Chinmay and Kaltenbach, Sebastian and Albergo, Michael S. and Du, Yilun and Menze, Bjoern},
  journal={arXiv preprint arXiv:2603.23398},
  year={2026}
}