ksusu/QHFlow2-rMD17

QHFlow2 — rMD17 Pre-trained Checkpoints

Pre-trained checkpoints for QHFlow2 on the revised MD17 (rMD17) dataset (DFT Hamiltonian prediction).

Paper: High-order Equivariant Flow Matching for Density Functional Theory Hamiltonian Prediction Authors: Seongsu Kim, Nayoung Kim, Dongwoo Kim, Sungsoo Ahn (KAIST) Venue: NeurIPS 2025 Code: github.com/seongsukim-ml/QHFlow2

Model Variants

Size	hidden_size	num_gnn_layers	Checkpoint Size
small_v2	64	3	313 MB
middle	128	3	975 MB

Molecules

Molecule	Atoms	Formula
aspirin	21	C₉H₈O₄
naphthalene	18	C₁₀H₈
salicylic_acid	16	C₇H₆O₃

File Structure

rmd-{molecule}/
  QHFlow_so2_v5_1_{size}_b10-rmd-{molecule}/
    weights-epoch=79-val_loss=0.0000000.ckpt

Quick Start

1. Install QHFlow2

git clone https://github.com/seongsukim-ml/QHFlow2.git
cd QHFlow2
pip install -e ".[fairchem]"

2. Download Checkpoints

# Download a single checkpoint
huggingface-cli download ksusu/QHFlow2-rMD17 \
  "rmd-aspirin/QHFlow_so2_v5_1_small_v2_b10-rmd-aspirin/weights-epoch=79-val_loss=0.0000000.ckpt" \
  --local-dir ckpt/rmd17

# Download all checkpoints (~3.9 GB)
huggingface-cli download ksusu/QHFlow2-rMD17 --local-dir ckpt/rmd17

Or in Python:

from huggingface_hub import hf_hub_download

path = hf_hub_download(
    repo_id="ksusu/QHFlow2-rMD17",
    filename="rmd-aspirin/QHFlow_so2_v5_1_small_v2_b10-rmd-aspirin/weights-epoch=79-val_loss=0.0000000.ckpt",
)

3. Download Dataset

Download from Google Drive and place under dataset/:

QHFlow2/dataset/
  rmd-aspirin/
  rmd-naphthalene/
  rmd-salicylic_acid/

4. Run Prediction

python -m qhflow2.experiment.train_md17 \
  mode=predict \
  dataset=rmd-aspirin \
  ckpt=ckpt/rmd17/rmd-aspirin/QHFlow_so2_v5_1_small_v2_b10-rmd-aspirin/weights-epoch=79-val_loss=0.0000000.ckpt

5. Python API

import torch
from qhflow2.models import get_model, get_default_model_args

args = get_default_model_args("md17")
args["version"] = "QHFlow_so2_v5_1"
args["hidden_size"] = 64
args["num_gnn_layers"] = 3
model = get_model(args)

ckpt = torch.load("weights-epoch=79-val_loss=0.0000000.ckpt", map_location="cpu")
state_dict = {k.replace("model.", ""): v for k, v in ckpt["state_dict"].items()}
model.load_state_dict(state_dict, strict=False)
model.eval()

Citation

@inproceedings{kim2025high,
  title={High-order Equivariant Flow Matching for Density Functional Theory Hamiltonian Prediction},
  author={Kim, Seongsu and Kim, Nayoung and Kim, Dongwoo and Ahn, Sungsoo},
  booktitle={Advances in Neural Information Processing Systems},
  year={2025}
}

License

MIT