ksusu/QHFlow2-QH9

QHFlow2 — QH9 Pre-trained Checkpoints

Pre-trained checkpoints for QHFlow2 on the QH9 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	Params	Checkpoint Size
small	64	3	~35M	139 MB
middle	128	3	~125M	494 MB
large	256	4	~530M	2.0 GB
extra_large	512	4	~2.8B	11 GB

Dataset Splits

Split	Description	Epochs
QH9Stable-random	Random train/test split	78
QH9Stable-size_ood	Size out-of-distribution	78
QH9Dynamic-300k-geometry	Geometry OOD (300k conformers)	33
QH9Dynamic-300k-mol	Molecular OOD (300k conformers)	33

File Structure

{split}/
  QHFlow_so2_v5_1_{size}-{split}/
    weights-epoch=XX-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

pip install huggingface_hub

# Download a single checkpoint
huggingface-cli download ksusu/QHFlow2-QH9 \
  "QH9Stable-random/QHFlow_so2_v5_1_small-QH9Stable-random/weights-epoch=78-val_loss=0.0000000.ckpt" \
  --local-dir ckpt/QH9

# Download all checkpoints for a specific split
huggingface-cli download ksusu/QHFlow2-QH9 \
  --include "QH9Stable-random/*" \
  --local-dir ckpt/QH9

# Download everything (~48 GB)
huggingface-cli download ksusu/QHFlow2-QH9 --local-dir ckpt/QH9

Or in Python:

from huggingface_hub import hf_hub_download

path = hf_hub_download(
    repo_id="ksusu/QHFlow2-QH9",
    filename="QH9Stable-random/QHFlow_so2_v5_1_small-QH9Stable-random/weights-epoch=78-val_loss=0.0000000.ckpt",
)

3. Run Prediction

cd QHFlow2

# Predict on QH9Stable test set
python -m qhflow2.experiment.train_qh9 \
  mode=predict \
  dataset=QH9Stable \
  dataset.split=random \
  model=QHFlow_so2_v5_1_small \
  ckpt=ckpt/QH9/QH9Stable-random/QHFlow_so2_v5_1_small-QH9Stable-random/weights-epoch=78-val_loss=0.0000000.ckpt

4. Python API

import torch
from qhflow2.models import get_model, get_default_model_args

# Build model
args = get_default_model_args("qh9")
args["version"] = "QHFlow_so2_v5_1"
args["hidden_size"] = 64         # small
args["num_gnn_layers"] = 3
model = get_model(args)

# Load checkpoint
ckpt = torch.load("weights-epoch=78-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()

5. Load via Lightning (full pipeline)

from qhflow2.pl_module import _get_model_by_pl_type
from omegaconf import OmegaConf

conf = OmegaConf.load("configs/qh9/config_flow_v2_simple.yaml")
LitModel = _get_model_by_pl_type("flow")
lit_model = LitModel.load_from_checkpoint(
    "weights-epoch=78-val_loss=0.0000000.ckpt",
    conf=conf,
    strict=False,
)
lit_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