BOA — Basis Overlap Architecture
Trained checkpoints for the ICLR 2026 paper A Function-Centric Graph Neural Network Approach For Predicting Electron Densities.
BOA is an equivariant graph neural network that predicts ground-state electron densities. Its message passing uses the overlap matrix of the basis functions that represent the predicted density, rather than treating the basis coefficients as generic node features.
Code: https://github.com/sciai-lab/boa
Checkpoints
| file | dataset | NMAE [%] (this checkpoint) | NMAE [%] (mean ± standard error) | seeds |
|---|---|---|---|---|
qm9_pyscf_large.ckpt |
QM9 (PySCF) | 0.106 | 0.116 ± 0.006 | 5 |
qm9_pyscf_small.ckpt |
QM9 (PySCF) | 0.113 | 0.13 ± 0.01 | 3 |
qm9_vasp_large.ckpt |
QM9 (VASP) | 0.132 | 0.1339 ± 0.0005 | 5 |
qm9_vasp_small.ckpt |
QM9 (VASP) | 0.137 | 0.1381 ± 0.0003 | 3 |
qm9_pyscf_small_small_cutoff.ckpt |
QM9 (PySCF) | 0.121 | - | 1 |
benzene.ckpt |
MD | 0.355 | 0.361 ± 0.003 | 3 |
resorcinol.ckpt |
MD | 0.362 | 0.371 ± 0.004 | 3 |
phenol.ckpt |
MD | 0.494 | 0.56 ± 0.03 | 3 |
malonaldehyde.ckpt |
MD | 0.585 | 0.61 ± 0.01 | 3 |
ethanol.ckpt |
MD | 0.705 | 0.710 ± 0.004 | 3 |
ethane.ckpt |
MD | 0.767 | 0.772 ± 0.002 | 3 |
These are the best seeds The paper reports the mean over multiple seeds for each dataset; each checkpoint here is the single best-performing seed, so its NMAE is better than the published figure.
The reduced-cutoff model (extrapolation)
qm9_pyscf_small_small_cutoff.ckpt is the model from §3.2 of the paper, trained
for extrapolation to molecules far larger than those seen in training. It is
the small configuration on QM9/PySCF with two reduced radii: message passing
6 Å → 3 Å and edge features 3 Å → 2 Å (config
configs/experiment/qm9_pyscf_small_small_cutoff.yaml).
Usage
Install the code from https://github.com/sciai-lab/boa, then:
from boa.model.module import ChgLightningModule
model = ChgLightningModule.load_from_checkpoint("qm9_pyscf_large.ckpt", map_location="cpu")
model.eval()
On the weights. BOA is trained with an exponential moving average, and all
reported numbers were measured with the EMA weights rather than the raw ones. In
these released files the state_dict is the EMA weights, so a plain load
gives you the evaluated model — no extra step. (Training checkpoints are not
like this: there state_dict holds the raw weights and ema.copy_to must be
applied first. If you evaluate one of these files with the repository's
boa/test.py, its ema.copy_to call is a no-op and remains correct.)