ksusu/QHFlow2-MD17

QHFlow2 — MD17 Pre-trained Checkpoints

Pre-trained checkpoints for QHFlow2 on the MD17 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
ethanol	9	C₂H₆O
malondialdehyde	9	C₃H₄O₂
uracil	12	C₄H₄N₂O₂

File Structure

{molecule}/
  QHFlow_so2_v5_1_{size}_b10-{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-MD17 \
  "ethanol/QHFlow_so2_v5_1_small_v2_b10-ethanol/weights-epoch=79-val_loss=0.0000000.ckpt" \
  --local-dir ckpt/md17

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

Or in Python:

from huggingface_hub import hf_hub_download

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

3. Download Dataset

Download from Google Drive and place under dataset/:

QHFlow2/dataset/
  ethanol/
  malondialdehyde/
  uracil/

4. Run Prediction

python -m qhflow2.experiment.train_md17 \
  mode=predict \
  dataset=ethanol \
  ckpt=ckpt/md17/ethanol/QHFlow_so2_v5_1_small_v2_b10-ethanol/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