Physics-Aligned Electronic Ground-State Learning Improves Generalization

Eike S. Eberhard; Xaver Kainz; Viktor Kotsev; Abdulrahman Aldossary; Stephan Günnemann

cs.LGOct 7, 2026

Physics-Aligned Electronic Ground-State Learning Improves Generalization

Authors: Eike S. Eberhard, Xaver Kainz, Viktor Kotsev, Abdulrahman Aldossary, Stephan Günnemann

Organizations: Technical University of Munich (TUM) · Munich Data Science Institute (MDSI) · Munich Center for Machine Learning (MCML) · NVIDIA

Abstract

Machine-learned interatomic potentials (MLIPs) excel at in-distribution tasks, accelerating drug and material development, yet they struggle to generalize out-of-distribution. We propose to push the cost-accuracy Pareto frontier by designing observable-agnostic electronic ground-state descriptor models (GSMs) with computational costs situated between MLIPs and Kohn-Sham density functional theory (KS-DFT). We align the learning objectives and architectures of GSMs with the governing equations of KS-DFT by enforcing physical constraints and removing optimization pressure on unphysical or irrelevant degrees of freedom. In our size-extrapolation experiments from QM9 to QM40, our combined contributions OrthoNormal-Loss (ON-Loss) and Grassmann Restricted Occupied-Orbital Training (GROOT) reach a 79.1% energy and 83.4% force mean absolute error (MAE) reduction over previous state-of-the-art density GSMs. For Hamiltonian GSMs, ON-Loss and Residual Optimal-gauge Conditioning-aware KS-Eq. Training (ROCKET) together reduce the energy and force MAEs of the strongest baseline by 99.8% and 95.9%, respectively. Using a self-consistency rejection criterion, we filter out extrapolation errors on QMugs, rejecting fewer than 0.4% of predictions while reaching an energy MAE of 0.07 mHa. Finally, we demonstrate the efficiency of label-free self-consistency fine-tuning, and transfer GSMs to reactive chemistry in Transition1x, reaching energy errors below chemical accuracy.

Figures & tables

Figure 1: GSM size extrapolation from QM9 to molecules up to $4\times$ larger QM40 molecules. The GROOT models generalize robustly, reaching errors below half chemical accuracy ( $\sim 0.8$ mHa) on 98% of the tested molecules. All observables are computed at a budget of exactly one Fock build averaged over 3 training seeds each. The GROOT variants are the best performing models for almost all observables except for the Kohn–Sham (self-consistency) energy where ROCKET is best. NequIP and MACE trained on forces still outperform even the best GSMs on forces, but physically aligned GSMs achieve an order of magnitude better energies compared to all trained MLIPs.

Observable	Leading order in $\\|\cdot\\|_{F}$
Unconstrained $\Delta\mathbf{P}$
$E_{\text{KS}}$	$\propto$	$\lambda_{\max}\\|\Delta\mathbf{P}\\|_{F}$
$\mu_{\rho}/L_{\rho}^{1}/L_{\rho}^{2}$	$\propto$	$\lambda_{\max}\\|\Delta\mathbf{P}\\|_{F}$
Manifold Constrained $\Delta\mathbf{P}$
$E_{\text{KS}}[\rho]$	$\propto$	$\lambda_{\max}^{2}\\|\Delta\mathbf{P}\\|_{F}^{2}$
$\mu_{\rho}/L_{\rho}^{1}/L_{\rho}^{2}$	$\propto$	$\lambda_{\max}\\|\Delta\mathbf{P}\\|_{F}$

Table 1: Worst-case leading-order observable errors in terms of matrix errors.

\lambda_{\min}

and

\lambda_{\max}

are the extreme eigenvalues of the overlap matrix

\mathbf{S}

for a given molecule, and

\Delta\epsilon_{\mathrm{HL}}

its HOMO-LUMO gap. The overlap factors are the worst-case conversion factors from non-orthonormal to orthonormal matrix errors, attained along the extreme overlap eigenvectors. Derivations are in App. C .

Figure 2: ON-Loss is "observable-isotropic" and improves size extrapolation. (a,b,c,d) Values log-binned by

\sqrt{\lambda_{i}\lambda_{j}}

, matching the directions

\bm{\Pi}_{ij}

of QM40 molecules ( table 5 ); shading indicates the standard deviation, histograms show the direction-density

\varrho

. (a,b) Mean-normalized relative sensitivity for the

L^{1}

density, dipole

\mu

, and energy

E

errors (line styles): (a)

\mathcal{L}_{\mathrm{NON}}

falls as

1/\raisebox{-0.4pt}[0.0pt][0.0pt]{

\sqrt{\lambda_{i}\lambda_{j}}

}

, whereas (b)

\mathcal{L}_{\mathrm{ON}}

stays approximately constant for all three observables. (c,d) Matrix error

\Delta\bar{x}_{b}\equiv\sqrt{\langle\mathrm{tr}(\bm{\Pi}_{ij}\Delta\mathbf{X})^{2}\rangle_{b}}

of the trained network, with

P_{\textrm{MINAO}}

for reference: (c)

\mathcal{L}_{\mathrm{NON}}

leaves error that follows

\varrho

rather than observable impact, while (d)

\mathcal{L}_{\mathrm{ON}}

concentrates it in the high-impact directions. (e–h) Size extrapolation experiments from QM9 to QM40 , featuring the same observables MAEs plus the force

F

acting on the atoms plotted against the electron count

N_{\mathrm{elec}}

. We compare GSMs under both

\mathcal{L}_{\mathrm{NON}}

and

\mathcal{L}_{\mathrm{ON}}

loss functions, as well as Hamiltonian variants with a training-set-optimal constant gauge shift

c_{t}

. Markers are means and bands the minimum and maximum across identical runs; every

\mathcal{L}_{\mathrm{ON}}

model beats its

\mathcal{L}_{\mathrm{NON}}

counterpart.

Figure 3: GROOT and ROCKET show robust extrapolation, and can be equipped with GSMs-specific error indicators ( table 28 ). GSMs and MLIPs trained on QMugs molecules with at most 150 electrons and evaluated on systems with up to 660 electrons. The table shows the fraction of molecules rejected by the self-consistency criterion and the observable errors of the remaining molecules, with the MAE including the rejected predictions in brackets below. Compared to linear scaling MLIPs like PaiNN ( Schütt et al., 2021 ) , NequIP ( Batzner et al., 2022 ) , and MACE ( Batatia et al., 2022 ) trained on energies only (meaning at equal reference compute budget), GROOT proj achieves an order of magnitude better energies, rejecting less than 1% of its own predictions, characterizing the achievable cost-accuracy trade-off of GSMs compared to MLIPs. The ERIC ( Eberhard et al., 2026a ) relative to MINAO shows that our density models yield additional SCF acceleration over the baselines. Observable MAEs vs. electron count. GROOT and ROCKET stay orders of magnitude below the $\mathcal{L}_{\mathrm{NON}}$ baselines with the margin widening as the system size increases.

Figure 4: Beyond size extrapolation: generalization across conformations and label-free adaptation to new chemistry. QMugs-trained models generalize to the $\beta=120^{\circ}$ cut of the 3BPA backbone-dihedral energy surface (left panel, table 29 ): the ground truth $(\alpha,\gamma)$ energy-surface, and the energy deviations of GSMs, with median (MAE-optimal) constant offset removed per model. Label-free self-consistency training (SCT) on Transition1x. The upper-right panel shows the validation loss against the self-consistency residual over 0, 2, 4, and 6 epochs of cosine-annealed SCT, suggesting that shorter training budgets might be feasible. The observable MAEs are resolved along the reaction coordinate $\xi$ , with the transition state at $\xi=0.5$ . The energies of the finetuned GSMs consistently achieve chemical accuracy (solid red line) on the accepted predictions, with GROOT proj and ROCKET rejecting the fewest predictions ( $0.2\%$ , table 30 ).

Appendix figures & tables46 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Many coefficient matrices, one density. (a) The smallest case,

B=3

and

O=2

: the occupied orbitals

c_{1},c_{2}

span a plane which determines the density. Rotating the arrows within the plane by

\mathbf{U}_{OO}

changes the orbitals, but not the plane. (b) For arbitrary

O,B\in\mathbb{N}

with

O<B

. The coefficients

\mathbf{C}\in\mathbb{R}^{B\times B}

split into an occupied block

\mathbf{C}_{O}\in\mathbb{R}^{B\times O}

and a virtual block

\mathbf{C}_{\mathrm{V}}\in\mathbb{R}^{B\times V}

. Only the occupied block determines the density matrix as seen in eq. 2 . Multiplication from the right by a block-diagonal orthogonal matrix

\mathbf{U}

does not change the density matrix, see eq. 5 .

Figure 6: AO versus Löwdin spatial decay for QM40 molecule 68311 at B3LYP/def2-SVP. Columns are $\mathbf{H}$ , $\mathbf{P}$ , $\mathbf{P}-\mathbf{P}_{\textrm{MINAO}}$ , and $\mathbf{H}-\mathbf{H}_{\textrm{SAP}}-c\mathbf{S}$ , each in the AO (a.k.a. NON) basis and after Löwdin orthogonalization, scaled by the dataset root-mean-square (RMS) $\sigma_{0}$ . AO entries decay with atom-pair distance. Löwdin entries remain large after the AO entries have dropped.

Parameter	Value	Notes
Basis set	def2-SVP	Ahlrichs split-valence double-zeta basis with polarization functions ( Weigend and Ahlrichs, 2005 ) .
XC functional	B3LYP	Three-parameter hybrid combining exact exchange ( Becke, 1993 ) with LYP ( Lee et al., 1988 ) and VWN local correlation ( Vosko et al., 1980 ) in the parameterization of Stephens et al. (1994) ; we use the VWN-RPA form, as Gaussian does.
Density fitting	RI-J, RI-K	Density fitting of both the Coulomb ( Vahtras et al., 1993 ) and the exact-exchange term ( Weigend, 2002 ) .
Auxiliary basis	def2-universal-JKFIT	One auxiliary basis serves both the RI-J and the RI-K contractions ( Weigend, 2008 ) .
Spin treatment	Restricted	Closed-shell restricted Kohn–Sham.
Quadrature grid	Level 3	PySCF grid level for the XC energy and potential ( Becke, 1988 ; Sun et al., 2020 ) .

Appendix

Table 2: DFT reference pipeline (GPU4PySCF 1.8.1, B3LYP/def2-SVP, RI-JK). This pipeline produces the evaluation references for all five datasets, with the sample counts of each split given in tables 4 , 5 , 6 , 7 and 8 . It also produces the training labels of the QMugs-trained checkpoints. The QM9 training labels come from a separate store computed with identical settings except for a looser convergence threshold ( table 4 ). Transition1x self-consistency training uses no labels.

Stage	Cost	Computed quantities
Cheap	One Fock build	The predicted density matrix $\hat{\mathbf{P}}$ is the model output for density readouts and the density matrix after diagonalization of the predicted $\hat{\mathbf{H}}$ ( eq. 1 ) for Hamiltonian readouts. From $\hat{\mathbf{P}}$ we build $\mathbf{F}[\hat{\mathbf{P}}]$ once. Every quantity that needs no further Fock build and no nuclear gradient counts as cheap.
SCF	Up to $2\times$ reference	A full SCF started from $\hat{\mathbf{P}}$ with the settings of table 2 . A run is counted as not converged after twice the cycle count of the MINAO-started reference SCF.
Forces	One nuclear gradient	Nuclear forces evaluated at the density matrix $\mathbf{P}(\mathbf{F}[\hat{\mathbf{P}}])$ , compared with the reference forces $-\partial E/\partial\mathbf{R}$ .

Appendix

Table 3: Compute stages, costs, and evaluated quantities for model predictions. Each split in tables 4 , 5 , 6 , 7 and 8 states the sample count evaluated at each stage. All errors are measured relative to the reference calculations of table 2 .

Parameter	Value	Notes
Release	$133{,}885$ molecules	Small organic molecules with up to 9 heavy atoms (C, N, O, F), enumerated from GDB-17, with equilibrium geometries relaxed at B3LYP/6-31G(2df,p).
Our copy	$130{,}831$ molecules	We drop the $3{,}054$ molecules that the authors flag as uncharacterized.
Geometries	0 K minima	Only the released thermochemistry refers to 298.15 K.
Elements	H, C, N, O, F	$1{,}923$ molecules contain F.
Size range		3–29 atoms, 1–9 heavy atoms, 10–74 electrons, 24–226 basis functions.
Training labels	$10^{-9}\,E_{h}$	The QM9 checkpoints are trained on labels converged to $10^{-9}\,E_{h}$ and otherwise identical to table 2 . The total energies agree with the $10^{-10}\,E_{h}$ evaluation reference to a median of $2.5\times 10^{-10}\,E_{h}$ , with a maximum of $1.5\times 10^{-5}\,E_{h}$ .

Appendix

Table 4: QM9 training splits and evaluation sample counts ( Ramakrishnan et al., 2014 ) . The split is a deterministic size cut on the total atom count including hydrogens and involves no random draw. Compute stages are defined in table 3 . The electron-count distribution is shown in fig. 7 .

Figure 7: Electron-count histogram of QM9. The vertical axis is the molecule count on a logarithmic scale. QM9 contains only H, C, N, O, and F.

Parameter	Value	Notes
Release	$162{,}954$ molecules	Drug-like molecules from ZINC14 with 10–40 heavy atoms (C, N, O, S, F, Cl), covering the size range of 88% of FDA-approved drugs. Geometries are pre-optimized with GFN2-xTB and relaxed at B3LYP/6-31G(2df,p), the level of theory of QM9; molecules with imaginary frequencies are removed.
Geometries	0 K minima
Element filter	$99{,}303$ molecules	H, C, N, O, F only, 78–272 electrons. No molecule with 20 or 21 heavy atoms passes the filter, since all of them contain S or Cl.
Electron bins	$[80,280)$	Ten bins of 20 electrons. The two molecules with 78 electrons fall below the first bin. The sparsest bin, 80–100 electrons, holds 311 molecules.
Validation (hyperparameter selection)
Rule	50 regular, 25 hard	Used to select learning rate and weight decay (App. L ). We remove them from the pool before drawing the test sets.

Appendix

Table 5: QM40 size-extrapolation dataset splits and evaluation sample counts ( Madushanka et al., 2024 ) . These checkpoints never see QM40 during training. To match their element vocabulary, we remove every molecule containing S or Cl. Compute stages are defined in table 3 . The electron-count distribution is shown in fig. 8 .

Figure 8: Electron-count histogram of QM40. The vertical axis is the molecule count on a logarithmic scale. The two series are the full dataset and its H, C, N, O, F subset.

Parameter	Value	Notes
Release	$1{,}992{,}984$ conformers of $665{,}911$ molecules	Neutral, closed-shell drug-like molecules from ChEMBL 27 with 3–100 heavy atoms (H, C, N, O, F, P, S, Cl, Br, I). Each molecule has three conformers, selected by clustering GFN2-xTB metadynamics snapshots and then optimized with GFN2-xTB.
Geometries	0 K minima	GFN2-xTB local minima; 300 K is only the temperature of the conformer search.
Our copy	$657{,}142$ molecules	All conformers via OpenQDC, 22–850 electrons. Molecules are identified by SMILES, which merges some ChEMBL entries.
Training split
Rule	$N_{\mathrm{elec}}\leq 150$	Every conformer with at most 150 electrons ( $184{,}583$ conformers of $60{,}367$ molecules; 3–22 heavy atoms, 4–54 atoms) that is not in the validation split; all conformers kept.
Elements	H, C, N, O, F, P, S, Cl	Molecules with Br or I are excluded.

Appendix

Table 6: QMugs dataset splits and evaluation sample counts for pre-training and size extrapolation ( Isert et al., 2021 ) . Every sample is a single conformer, and every split is disjoint by molecule. Sets with one conformer per molecule fix that choice in a single draw (seed 0). Compute stages are defined in table 3 . The electron-count distribution is shown in fig. 9 .

Figure 9: Electron-count histogram of QMugs. The vertical axis is the molecule count on a logarithmic scale. The two series are the full set and the H, C, N, O, F subset ( Isert et al., 2021 ) .

Parameter	Value	Notes
Release	$13{,}993$ frames	A single flexible molecule, 3-(benzyloxy)pyridin-2-amine, with three consecutive rotatable bonds. The MD frames come from 25 ps Langevin trajectories at 300, 600 and 1200 K (1 fs time step) started in five conformational pockets; the dihedral scans are constrained DFT geometry optimizations. The released labels are $\omega$ B97X/6-31G(d), computed with ORCA.
Molecule	C 12 H 12 N 2 O	27 atoms, 106 electrons, 270 basis functions.
Used subsets	Dihedral scans	$7{,}047$ frames. We do not use the two training sets (500 frames each) or the MD test sets ( $1{,}669$ / $2{,}138$ / $2{,}139$ frames at 300 / 600 / 1200 K).
Dihedrals	$\alpha,\beta,\gamma$	Three consecutive rotatable bonds of the benzyl-ether linker: $\alpha$ about the pyridine C–O bond, $\beta$ about the O–CH 2 bond, and $\gamma$ about the CH 2 –phenyl bond. A fourth constraint fixes the amino group.
Grid	$47\times 50$	$\alpha=55^{\circ},\dots,239^{\circ}$ in steps of $4^{\circ}$ ; $\gamma=5^{\circ},\dots,299^{\circ}$ in steps of $6^{\circ}$ .
Geometries	0 K	Constrained minima at each $(\alpha,\gamma)$ with $\beta$ fixed.

Appendix

Table 7: 3BPA conformational transfer dataset and scan sample counts ( Kovács et al., 2021 ) . These checkpoints never see 3BPA during training. We score each fixed-dihedral scan in full, without subsampling, because a subset of a potential-energy cut would be a different curve. Compute stages are defined in table 3 .

Parameter	Value	Notes
Release	$9{,}654{,}813$ frames	Configurations on and around the reaction paths of $10{,}073$ reactions from Grambow et al. (2020) , each with up to seven heavy atoms (C, N, O). The frames are images of nudged-elastic-band and climbing-image NEB runs, labeled at $\omega$ B97x/6-31G(d).
Elements	H, C, N, O	4–23 atoms.
Geometries	Reaction paths	Non-equilibrium configurations: many lie on or near reaction barriers.
Release split	$9{,}561$ / 225 / 287 reactions	Training / validation / test reactions of the release.
Reaction coordinate	$\xi\in[0,1]$	$\xi=\tfrac{1}{2}(1-d_{\mathrm{TS}}/d_{\mathrm{R,TS}})$ on the reactant side and $\xi=\tfrac{1}{2}(1+d_{\mathrm{TS}}/d_{\mathrm{P,TS}})$ on the product side, where $d$ is the Kabsch RMSD; R, TS and P sit at $\xi=0$ , $0.5$ and $1$ .
Training split (label-free)

Appendix

Table 8: Transition1x reactive chemistry dataset splits and sample counts for self-consistency training ( Schreiner et al., 2022 ) . We inherit the release split, which assigns each molecular formula to one split, so our training, validation and test sets share no formula and no reaction. Within each reaction we keep the reactant (R), transition state (TS) and product (P) and a window of path images around the TS. We sign each image’s Kabsch root-mean-square deviation (RMSD) to the TS by the side of the TS it lies on. We divide

[-0.5,0.5]

Å of this signed distance into 20 bins of 0.05 Å and keep the image nearest each bin center, dropping empty bins and geometric duplicates. Compute stages are defined in table 3 .

model type	supervision target	$\sigma_{0}[\text{e}/\text{Ha}/-]$	$1/\sigma_{0}$
$\Delta\mathbf{P}_{\textrm{MINAO}}$	$\mathbf{S}^{1/2}\!\left(\mathbf{P}^{\mathrm{ref}}-\mathbf{P}^{\textsc{minao}}\right)\!\mathbf{S}^{1/2}$	$0.0350$	$28.6$
$\Delta\mathbf{P}_{\textrm{MINAO}}$ (AO)	$\mathbf{P}^{\mathrm{ref}}-\mathbf{P}^{\textsc{minao}}$	$0.0295$	$33.9$
$\Delta\mathbf{H}_{\textrm{SAP}}$	$\mathbf{S}^{-1/2}\!\left(\mathbf{H}^{\mathrm{ref}}-\mathbf{H}^{\textsc{sap}}\right)\!\mathbf{S}^{-1/2}$	$0.0349$	$28.7$
$\Delta\mathbf{H}_{\textrm{SAP}}$ (AO)	$\mathbf{H}^{\mathrm{ref}}-\mathbf{H}^{\textsc{sap}}$	$0.0471$	$21.2$
$\Delta\mathbf{H}_{\textrm{SAP}}$ fitted gauge	$\mathbf{S}^{-1/2}\!\left(\mathbf{H}^{\mathrm{ref}}-\mathbf{H}^{\textsc{sap}}-\bar{c}\,\mathbf{S}\right)\!\mathbf{S}^{-1/2}$	$0.0148$	$67.6$
$\Delta\mathbf{H}_{\textrm{SAP}}$ fitted gauge (AO)	$\mathbf{H}^{\mathrm{ref}}-\mathbf{H}^{\textsc{sap}}-\bar{c}_{\mathrm{AO}}\,\mathbf{S}$	$0.0171$	$58.5$

Appendix

Table 9: Output-scale normalization constants for different model types.

\sigma_{0}

is the element-wise root-mean-square of the supervision target in the representation the readout emits into, using the Löwdin frame for the Löwdin-basis heads and the AO basis for the AO-basis heads. We fit all constants on the full training split (

103{,}790

molecules) and verify them end-to-end through the training data pipeline on

20

random training molecules with agreement within

2.7\%

. The readout applies the constant internally. The additive heads compose

\mathrm{init}+\sigma_{0}\,\mathrm{sym}(\hat{\mathbf{M}})

, while the rotation head scales the decoder output before generator construction (

\hat{\mathbf{X}}

is linear in

\hat{\mathbf{M}}

), keeping the network output at unit scale while losses and gauge constants retain their physical units.

\bar{c}=0.404129

Ha is the training-set constant gauge.

\bar{c}_{\mathrm{AO}}=0.382862

Ha is the corresponding constant optimal in the AO Frobenius metric. The two metrics give different constants with spreads of

37.3

and

41.0

mHa. Because

\sigma_{0}

anticorrelates with system size (

r=-0.86

against the basis dimension), constants must be fitted on size-representative samples. For the rotation head, the decoder output

\hat{\mathbf{M}}

enters through the projector sandwich

\hat{\mathbf{X}}(\hat{\mathbf{M}})

( eq. 34 ), which pins only the occupied–virtual block of

\hat{\mathbf{M}}

at unit weight (

\mathbf{P}_{0}\mathbf{S}

is idempotent) and annihilates the rest. The value

0.0110

is the first-order pinned occupied–virtual block rms, measured against the purified minao projector. The true Grassmann-log value runs about

7\%

higher (

\approx 0.0117

), and we use the first-order convention because the fit and the configuration must share one convention. Under the Grassmann interpretation,

\sigma_{0}

is a rotation-angle scale rather than a matrix-element scale. The gap-shifted target admits a closed-form constant because the shift lies in the span of

\{\mathbf{I},\bm{\Pi}^{\mathrm{ref}}\}

, giving

\sigma_{0}^{2}(\xi)=\sigma_{0,\Delta H}^{2}+2b_{1}\xi+b_{2}\xi^{2}

with

b_{1}=\mathbb{E}\!\left[\mathrm{tr}\,\Delta-\mathrm{tr}(\Delta\Pi)\right]/B^{2}=2.03\times 10^{-3}

and geometric

b_{2}=\mathbb{E}\!\left[B-N\right]/B^{2}=4.70\times 10^{-3}

. No re-measurement is needed for a new shift (

\sigma_{0}=0.0378/0.0409/0.0471

\xi=0.05/0.1/0.2\,\mathrm{Ha}

Hyperparameter	Value	Notes
Number of layers	4
Sphere channels	192	Shared by backbone and decoder.
Maximum degree $\ell_{\max}$	4	def2-SVP off-diagonal $2\otimes 2$ requires $\ell\leq 4$ .
Edge channels	128
Node hidden channels	128
Radial basis	128	fixed Gaussians on $[0,5]$ Å

Appendix

Table 10: QHFlow2 backbone ( Kim et al., 2026 ) . SO(2)-equivariant edgewise–nodewise updates in the style of eSEN ( Fu et al., 2025 ) , with the generative flow replaced by the GROOT readout.

	PaiNN	NequIP	MACE
Interaction layers	3	4	3
Features / channels	128	128 per $l$	256
Max. angular order	$l=1$ (vectors)	$l_{\max}=3$	$L_{\max}=2$
Correlation order	–	–	3 (body order 4)
Parity	–	off	–
Radial basis	Gaussian, 20	Bessel, 8	Bessel, 8

Appendix

Table 11: Architecture of the PaiNN, NequIP and MACE baselines in the heat maps. All three models sum per-atom energies over the molecule; the PaiNN and MACE MLPs use SiLU. MACE hidden features are

256{\times}(0e{+}1o{+}2e)

, and its edge spherical harmonics reach

l=3

. The NequIP per-element energy scale is the per-element force RMS of the QM9 training split for the E and E+F models, and 1 for QMugs, which has no forces. The NequIP per-element shift is a least-squares fit to the training energies. MACE uses linear read-outs after interaction layers 1 and 2 and an MLP after layer 3, summed over layers.

Figure 10: Validation-loss trajectories for learning-rate selection on QMugs. (a) Hamiltonian models. (b) Density models. Each panel reports the selected learning rate above the axes and colors the remaining rates by their position on the sweep. Solid lines show 24-epoch runs, dashed lines show 12-epoch brackets, and crosses mark runs that stopped early. Completed 24-epoch runs take 35–39 h on one NVIDIA H200 GPU, or about 1.5 h per epoch.

Hyperparameter	Value	Notes
Batch size	$16$
Optimizer	Muon	Muon on the 2D hidden weight matrices, AdamW on all remaining parameters. Both optimizer branches use Nesterov momentum, following the Optax (contrib) defaults ( Babuschkin et al., 2020 ) . We also tried the same momentum as used for standard Adam but observed no improvements.
Muon momentum	$0.95$	Default ( Jordan et al., 2024 ) .
Muon Newton–Schulz steps	$5$	Default ( Jordan et al., 2024 ) .
Adam $(\beta_{1},\beta_{2})$	$(0.9,\,0.999)$	Default ( Kingma and Ba, 2015 ) .
Muon consistent RMS	$0.2$	Each orthogonalized $m\times n$ Muon update is scaled by $0.2\sqrt{\max(m,n)}$ , giving approx. shape-independent update RMS ( Liu et al., 2025 ) .

Appendix

Table 12: Training and optimization hyperparameters. Muon optimizer configuration. The optimizer configuration is shared by all reported runs. Family-specific deviations are stated in the notes. Model-init and training-PRNG seeds are drawn from OS entropy per run and logged.

Hyperparameter	Value	Notes
Optimal shift std $\sigma_{c}$	$37.3\,\mathrm{mHa}$	Standard deviation of the $L_{2}$ -optimal gauge shifts over the training split.
Restraint weight $\alpha$	$1.161$	Calibrated as $\alpha_{\textrm{ROCKET}}=\mathbb{E}[\mathcal{L}_{\mathrm{min}}(\Delta\widetilde{\mathbf{H}})]/\sigma_{c}^{2}$ over the full training split ( $103{,}790$ molecules).
EMA decay	$0.99$	Decay of the EMA of the optimal overlap shift center. Initialized to training split optimal $0.404$ .

Appendix

Table 13: ROCKET restraint hyperparameters. The restraint weight

\alpha_{\textrm{ROCKET}}

is the

3\sigma

weight

\alpha/9

with

\alpha=\mathbb{E}[\mathcal{L}_{\mathrm{min}}]/\sigma_{c}^{2}

fitted on the full training split. A three-standard-deviation gauge offset then costs as much as the mean shift-minimized loss. See section 4 .

					QM40 loss		$\\|W\\|_{\mathrm{rms}}$
ep	lr	wd	seed	val loss	reg. 50	hard 25	Muon	Adam
abs $P$ -learning (no-rescale)
16	3e-4	1e-2	995756	5.73	16.3	29.6	0.16	0.55
16	1e-3	1e-2	995756	4.82	14.8	27.8	0.17	0.44
16	3e-3	1e-2	995756	4.48	13.8	26.0	0.26	0.40
"	"	"	684851	4.40	14.1	36.9	0.25	0.41

Appendix

Table 14: Validation and probe loss of density models trained with the NON loss across learning rates, weight decays, and seeds. Masked decay with

\sigma_{0}

output norm. Target is

\Delta

-MINAO except for the absolute-readout block. Validation and QM40 loss are the AO objective (masked Frobenius root-mean-square error, RMSE) on the QM9 validation split at the final epoch and on the fixed probe molecules (reg. 50 / hard 25).

\|W\|

columns are per-entry RMS of the final EMA weights by optimizer branch. ON-Loss models are in table 15 .

					QM40 loss		$\\|W\\|_{\mathrm{rms}}$
ep	lr	wd	seed	val loss	reg. 50	hard 25	Muon	Adam
ON-Loss
12	1e-3	1e-2	0	1.46	7.54	9.06	0.17	0.48
"	"	"	749811	1.47	7.10	9.12	0.17	0.47
12	1e-3	3e-2	0	1.40	7.99	9.51	0.10	0.35
12	3e-3	1e-3	0	1.43	6.03	8.60	0.42	0.59

Appendix

Table 15: Validation loss of density models trained with the Löwdin-frame ON-Loss. Masked decay with

\sigma_{0}

output norm. Target is

\Delta

-MINAO. Validation is the training objective (masked Frobenius RMSE, Löwdin picture) on the QM9 validation split at the final epoch. The QM40-loss columns are the same objective evaluated on the fixed probe molecules (reg. 50 / hard 25, same units). The density family ranks models (Spearman

+0.87/+0.79

E

\|W\|

columns are per-entry RMS of the final EMA weights by optimizer branch. NON-loss models are in table 14 . GROOT models are in table 16 . Prediction-frame and auxiliary-supervision models are in table 17 .

					QM40 loss		$\\|W\\|_{\mathrm{rms}}$
ep	lr	wd	seed	val loss	reg. 50	hard 25	Muon	Adam
+ GROOT proj
16	1e-3	3e-2	193360	0.38	1.24	2.48	0.09	0.32
16	3e-3	3e-3	193360	0.39	2.33	2.52	0.36	0.47
16	3e-3	1e-2	193360	0.39	1.17	2.97	0.22	0.36
"	"	"	141727	0.39	1.06	2.44	0.21	0.35

Appendix

Table 16: Validation and probe loss of GROOT density models trained with ON-Loss. Masked decay with

\sigma_{0}

output norm and a

\Delta

-MINAO target. Validation is the training objective on the QM9 validation split at the final epoch. The QM40-loss columns evaluate that objective on the fixed probe molecules (reg. 50 / hard 25).

\|W\|

columns are per-entry RMS of the final EMA weights by optimizer branch.

					QM40 loss		$\\|W\\|_{\mathrm{rms}}$
ep	lr	wd	seed	val loss	reg. 50	hard 25	Muon	Adam
Density GROOT proj , pre/post supervision
16	3e-3	1e-2	193360	(2.17)	3.75	5.77	0.23	0.38
Density GROOT proj , AO-frame prediction
16	3e-3	1e-2	193360	0.58	1.39	3.03	0.21	0.34
Hamiltonian ROCKET, AO-frame prediction

Appendix

Table 17: Prediction-frame and auxiliary-supervision ablations for density and Hamiltonian models under ON-Loss. Density target is

\Delta

-MINAO and Hamiltonian target is

\Delta

-SAP with the training-set optimal constant gauge. Parenthesized validation is the model’s own training objective when that objective differs from the table metric. The QM40-loss columns use the Löwdin objective on the fixed probe molecules (reg. 50 / hard 25). The Hamiltonian AO-frame run was killed before the QM40 probe.

					QM40 loss		$\\|W\\|_{\mathrm{rms}}$
ep	lr	wd	seed	QM9 val	reg. 50	hard 25	Muon	Adam
abs $H$ -learning
16	1e-4	1e-3	684464	0.78	3.95	12.2	0.18	0.62
$\Delta H_{\mathrm{SAP}}$ -learning
16	3e-6	1e-3	684464	2.12	8.7	18.4	0.18	0.63
16	1e-5	1e-3	684464	1.12	5.3	15.1	0.18	0.63

Appendix

Table 18: Validation loss of Hamiltonian models trained with the NON loss, with and without a constant gauge shift. Masked decay with

\sigma_{0}

output norm. Target is

\Delta

-SAP, with a training-set constant gauge in the second block only. Validation is the NON training objective on the QM9 validation split at the final epoch. The QM40-loss columns evaluate that objective on the fixed probe molecules (reg. 50 / hard 25). The covariant loss is gauge-dominated and nearly flat across models, quoted for completeness, not for ranking.

\|W\|

columns are per-entry RMS of the final EMA weights by optimizer branch. ON-Loss models are in table 20 .

					QM40 loss		$\\|W\\|_{\mathrm{rms}}$
ep	lr	wd	seed	QM9 val	reg. 50	hard 25	Muon	Adam
$\Delta H_{\mathrm{SAP}}$ -learning
16	1e-4	1e-3	684464	11.8	12.1	56	0.18	0.63
$\Delta(H_{\mathrm{SAP}}+c_{\textrm{train-set-fitted}}S)$ -learning
16	3e-6	1e-3	684464	36.0	31.6	87	0.18	0.63
16	1e-5	1e-3	684464	9.17	12.1	110	0.18	0.63

Appendix

Table 19: Validation and probe loss of Hamiltonian models trained with WALoss. Target is

\Delta

-SAP, with a training-set constant gauge in the second block only. Validation is the WALoss training objective on the QM9 validation split at the final epoch. The QM40-loss columns evaluate each run’s own WALoss objective on the fixed probe molecules (reg. 50 / hard 25).

\|W\|

columns are per-entry RMS of the final EMA weights by optimizer branch.

					QM40 loss		$\\|W\\|_{\mathrm{rms}}$
ep	lr	wd	seed	QM9 val	reg. 50	hard 25	Muon	Adam
$\Delta H_{\mathrm{SAP}}$ -learning
16	3e-5	3e-3	670079	1.60	3.6	9.7	0.18	0.63
"	"	"	435126	1.61	3.7	9.9	0.18	0.62
"	"	"	301732	1.63	3.7	10.0	0.18	0.63
$\Delta(H_{\mathrm{SAP}}+c_{\textrm{train-set-fitted}}S)$ -learning

Appendix

Table 20: Validation and probe loss of Hamiltonian models trained with ON-Loss. Target is

\Delta

-SAP, with a training-set constant gauge in the

\Delta(H_{\mathrm{SAP}}+cS)

block only. Validation is the training objective on the QM9 validation split at the final epoch. The QM40-loss columns evaluate that objective on the fixed probe molecules (reg. 50 / hard 25).

\|W\|

columns are per-entry RMS of the final EMA weights by optimizer branch.

					QM40 loss		$\\|W\\|_{\mathrm{rms}}$
ep	lr	wd	seed	QM9 val	reg. 50	hard 25	Muon	Adam
1st stage (unrestrained gauge)
16	1e-4	1e-2	443905	0.96	41.0	32.5	0.17	0.60
16	3e-4	1e-2	443905	0.75	14.2	24.5	0.16	0.56
16	1e-3	1e-2	443905	0.67	24.6	26.8	0.17	0.44
2nd stage (restrained gauge)

Appendix

Table 21: Validation and probe loss of Hamiltonian models trained with the ROCKET loss. Target is

\Delta

-SAP with a training-set constant gauge. Validation is the model’s training objective on the QM9 validation split at the final epoch. The QM40-loss columns evaluate that objective on the fixed probe molecules (reg. 50 / hard 25).

\|W\|

columns are per-entry RMS of the final EMA weights by optimizer branch.

LR	Validation self-consistency residual after epoch
	1 †	2 †	3	4 †	5	6	7	8	9	10	11	12 †
GROOT rot – density residual, $10^{-4}$ (dimensionless)
$3{\times}10^{-5}$ ∗	00 12.4	000 9.10	8.47	7.59	7.00	6.98
$10^{-4}$	000 9.08	000 7.41	7.00	6.65	6.20	5.33	5.26	4.98	4.88	4.84	4.81	4.81
$3{\times}10^{-4}$	000 8.13	000 7.42	6.71	5.97	5.86	5.25	5.01	4.66	4.34	4.28	4.23	4.19
$10^{-3}$	00 10.6	000 8.75	8.45	7.38	6.36	5.50	5.19	4.71	4.41	4.15	4.07	4.04

Appendix

Table 22: Validation loss across learning rates and epochs for Transition1x self-consistency training ( table 8 ). Entries report the mean training objective residual over the Transition1x validation split in units of

10^{-4}

, used to select checkpoints for table 30 . For ROCKET 3rd stage, values omit the gauge-restraint term. LR denotes the peak learning rate of the warmup-cosine schedule. Underlined values mark the lowest validation loss per configuration and epoch column, identifying the selected checkpoint. † Epochs evaluated in the Transition1x test evaluation table ( table 30 ). ∗ Runs stopped before configured completion. Each block reports a distinct model residual and cannot be compared across blocks.

					$E$ [mEh]		dipole [D]		$L_{1}$ [ $e$ ]
	$n$	lr	wd	loss	reg.	hard	reg.	hard	reg.	hard	cata.
NON losses
Density
abs. $P$	3	3e-3	1e-2	$14.3$	259	1 655	$0.76$	${\color[rgb]{1,0,0}31.9}$	${\color[rgb]{1,0,0}0.50}$	${\color[rgb]{1,0,0}4.13}$	11.7
$\Delta P_{\mathrm{MINAO}}$	3	3e-3	1e-2	$13.7$	76	263	$0.51$	${\color[rgb]{1,0,0}4.50}$	${\color[rgb]{1,0,0}0.38}$	${\color[rgb]{1,0,0}1.4}$	${\color[rgb]{1,0,0}1.3}$
Hamiltonian

Appendix

Table 23: QM40 probe observables of the production runs. Every model whose epoch, learning rate, and weight decay are set to its production value, aggregated over its entropy-seed replicates as the mean over seeds. The loss column evaluates each family’s own training objective on the fixed probe molecules (reg. 50,

\times 10^{4}

) and is comparable only within a block. The cata. column sums the reg. 50 and hard 25 catastrophe counts.

E

is the one-Fock-build electronic-energy MAE in mEh, and a catastrophe is

E>1\,000

mEh. Bold and underlined entries are the lowest and second-lowest values in each observable column, respectively. Red marks an error above

3

10

mEh for regular/hard

E

1

4

D for the dipole,

0.2

1\,e

for

L_{1}

, or cata.

>1

		$r$		$I$ (base 2)
Model	Seeds	NON	ON	NON	ON
$\|\Delta E\|$ [m $E_{h}$ ]
$\mathbf{P}_{\mathrm{base}}$ (NON)	3	$0.595$	$\mathbf{0.669}$	$0.38$	$\mathbf{0.56}$
$\mathbf{P}_{\mathrm{base}}$ (ON)	3	$0.778$	$\mathbf{0.909}$	$0.87$	$\mathbf{1.58}$
$\mathbf{H}_{\mathrm{base}}$ (NON)	1	$0.924$	$\mathbf{0.932}$	$1.17$	$\mathbf{1.32}$
$\mathbf{H}_{\mathrm{base}}$ (ON)	3	$0.555$	$\mathbf{0.802}$	$0.30$	$\mathbf{1.06}$

Appendix

Table 24: Correlation and mutual information between Frobenius loss and observable errors across QM40 plate panels. We report the Pearson correlation

r

between

\log_{10}

Frobenius loss and

\log_{10}

observable error, alongside mutual information

I

in bits estimated with

k

-nearest neighbors (

k=5

). Statistics are evaluated for each model pooled across seeds, corresponding to the scatter points in each plate panel. NON denotes the raw congruence

\|\Delta\mathbf{X}\|_{F}

, and ON denotes the Löwdin congruence

\|\mathbf{S}^{+1/2}\Delta\mathbf{X}\mathbf{S}^{+1/2}\|_{F}

for densities or

\|\mathbf{S}^{-1/2}\Delta\mathbf{X}\mathbf{S}^{-1/2}\|_{F}

for Hamiltonians. Bold text marks the superior value in each pair, where larger values indicate stronger predictive alignment. The evaluation uses 1000 QM40 molecules per seed without exclusions, with force rows restricted to the 200 molecules per seed that possess reference forces.

Figure 11: Energy error $E$ against Frobenius loss. Each point represents one QM40 test molecule and seed on logarithmic axes. Rows correspond to model families $P_{\mathrm{base}}$ , $H_{\mathrm{base}}$ , and $H_{\mathrm{base}}+c_{t}\mathbf{S}$ . Column blocks indicate the training objective ( $\mathcal{L}_{\mathrm{NON}}$ followed by $\mathcal{L}_{\mathrm{ON}}$ ). Within each block, the left panel evaluates the raw loss $\|\Delta\mathbf{X}\|_{F}$ and the right panel evaluates the Löwdin loss ( eq. 3 ). Individual panels report $r$ and $I$ , with Löwdin panels displaying the gain over the raw counterpart. Table 24 compiles summary statistics across all four observables.

Figure 12: Dipole error against Frobenius loss. Panels share the models, layout, pooling, and annotations of fig. 11 . The vertical axis displays the dipole error $\|\Delta\bm{\mu}\|$ in Debye.

Figure 13: Density $L^{1}$ error against Frobenius loss. Panels share the models, layout, pooling, and annotations of fig. 11 . The vertical axis displays the real-space density error $\|\Delta\rho\|_{L^{1}}$ in electrons.

Figure 14: Force error against Frobenius loss. Panels share the models, layout, pooling, and annotations of fig. 11 . The vertical axis displays the per-atom mean force error $\langle\|\Delta\mathbf{F}_{a}\|\rangle$ . Because only 200 molecules per seed carry reference forces, these panels contain fewer samples than the preceding figures.

Model	LR	WD	Excluded	Self-consistent	Dipole	$L_{1}$	Forces	HOMO–LUMO	ERIC
			(%)	energy (mHa)	(D)	(e)	(mHa/Å)	(mHa)	(ratio)
				$\|\hat{E}_{\mathrm{KS}}-E^{*}\|$	$\\|\Delta\bm{\mu}\\|_{2}$	$\\|\Delta\rho\\|_{L^{1}}$	$\langle\\|\Delta\mathbf{F}_{a}\\|_{2}\rangle_{a}$	$\|\Delta\varepsilon_{\mathrm{gap}}\|$
Density-space models
$\mathbf{P}_{\mathrm{abs.}}$ (NON)	$3{\times}10^{-3}$	$10^{-2}$	0 62.9 [62, 66]	0 9.65 000 0 (12.1) [9.1, 10] (12, 13)	0 0.0419 0 (0.0553) [0.040, 0.043] (0.055, 0.056)	0 0.0333 0 (0.0428) [0.032, 0.034] (0.043, 0.043)	0 0.974 0 0 (1.26) [0.94, 1.0] (1.3, 1.3)	0 0.578 0 0 (0.774) [0.56, 0.62] (0.77, 0.78)	0 0.745 (0.754) [0.74, 0.75] (0.75, 0.75)
$\mathbf{P}_{\mathrm{abs.}}$ (NON), no rescale	$3{\times}10^{-3}$	$10^{-2}$	0 61.8 [60, 64]	10.4 0000 0 (12.2) [10, 11] (12, 13)	0 0.0421 0 (0.0545) [0.042, 0.042] (0.054, 0.055)	0 0.0337 0 (0.0427) [0.034, 0.034] (0.042, 0.043)	0 1.00 00 0 (1.27) [1.00, 1.0] (1.3, 1.3)	0 0.570 0 0 (0.759) [0.56, 0.59] (0.75, 0.78)	0 0.745 (0.754) [0.74, 0.75] (0.75, 0.75)
$\mathbf{P}_{\mathrm{base}}$ (NON)	$3{\times}10^{-3}$	$10^{-2}$	0 21.2 [19, 25]	0 8.47 000 00 (9.32) [8.4, 8.5] (9.2, 9.6)	0 0.0512 0 (0.0563) [0.051, 0.052] (0.055, 0.057)	0 0.0399 0 (0.0439) [0.039, 0.040] (0.043, 0.045)	0 1.19 00 0 (1.31) [1.2, 1.2] (1.3, 1.3)	0 0.741 0 0 (0.813) [0.73, 0.77] (0.79, 0.86)	0 0.751 (0.755) [0.75, 0.75] (0.75, 0.76)

Appendix

Table 25: Observable errors on the QM9 test split for QM9-trained models ( table 4 ). Each cell reports the mean over seeds of the per-molecule mean absolute error, with the seed range in square brackets. Values in round brackets average all molecules of the lane before outlier exclusion, identical across models. Per column, the best bracketed mean and the best single seed are underlined. Predictions with DIIS residual

r_{\mathrm{DIIS}}/\sqrt{N_{\mathrm{elec}}}>30

mHa are excluded as outliers, where

r_{\mathrm{DIIS}}=\sqrt{8}\,\|\mathbf{S}^{-1/2}(\mathbf{F}\mathbf{P}\mathbf{S}-\mathbf{S}\mathbf{P}\mathbf{F})\mathbf{S}^{-1/2}\|_{F}

for scored density

\mathbf{P}

and Fock matrix

\mathbf{F}=\mathbf{F}[\mathbf{P}]

. Unbracketed values average kept molecules only, and Excluded reports the lost molecule percentage. Models losing all molecules display a dash before the bracketed value.

\hat{E}_{\mathrm{KS}}

is the KS-DFT energy evaluated at the predicted density ( section 5 ), and

E^{*}

is the converged reference energy. All observables evaluate 2000 molecules before outlier exclusion. ERIC reports SCF iterations relative to the MINAO guess, averaged over the full lane with unconverged runs capped at 2. LR and WD denote peak learning rate and weight decay from the training configuration. † Single seed without range.

Figure 15: Density-readout models on QM40, resolved over electron count. All models use the QM9-trained checkpoint of table 4 evaluated on the out-of-distribution QM40 set of table 5 , binned by the molecular electron count

N_{\mathrm{elec}}

and given a one-Fock-build budget. We compare the absolute density readout, its unrescaled variant, the

\Delta\mathbf{P}_{\mathrm{MINAO}}

readout in the AO basis, the same

\Delta

-readout in the Löwdin basis (ON-Loss), and the two manifold-restricted variants GROOT proj and GROOT rot ( section 4 ). Lines are means over the entropy-seed replicates stated in the legend and the shaded band spans the seed spread, so single-seed models carry no band; the vertical axes of all error panels are logarithmic. ERIC is the effective relative iteration count ( Eberhard et al., 2026a ) , where values below one mean fewer SCF iterations than the reference initial guess.

Figure 16: Density-readout models on QMugs, resolved over electron count. Same QM9-trained checkpoints, evaluation protocol, binning and aggregation as fig. 15 , on the larger QMugs molecules ( fig. 9 ), which reach

N_{\mathrm{elec}}\approx 650

against

\approx 270

on QM40. GROOT rot is the only model whose error panels stay flat over the whole range and it leads every error panel, by up to three orders of magnitude on the energies; together with the Löwdin

\Delta

-readout it is also the only model that stays below the reference guess in ERIC, whereas GROOT proj rises above it from

N_{\mathrm{elec}}\approx 320

on and climbs towards the cap on the largest molecules. The evaluation set is the 950-molecule qmugs_ebins200_cheap50 subset of table 27 .

Figure 17: Hamiltonian-readout models on QM40, resolved over electron count. Same evaluation protocol, binning and aggregation as fig. 15 , for the

\Delta\mathbf{H}_{\mathrm{SAP}}

readouts. The models span the representation (AO versus Löwdin basis, eq. 3 ), the loss (WALoss versus the ROCKET objective of eq. 9 ) and the gauge treatment (none, a training-set fitted constant shift, and the EMA-centered restraint of eq. 10 ), with the occupied-plus-LUMO masked variant as the final model. The Löwdin-basis ROCKET models separate from the AO-basis models by two to three orders of magnitude on every energy panel, and are the only models that stay below the reference guess in ERIC across the whole range.

Figure 18: Hamiltonian-readout models on QMugs, resolved over electron count. Same protocol as fig. 16 , for the

\Delta\mathbf{H}_{\mathrm{SAP}}

readouts. The AO-basis model sits one to two orders of magnitude above the Löwdin-basis ROCKET models on every panel and is pinned at the ERIC cap across the range. Among the ROCKET models the gauge-EMA, occupied-plus-LUMO masked variant leads every panel, but unlike on QM40 ( fig. 17 ) every Hamiltonian model exceeds the reference guess in ERIC beyond

N_{\mathrm{elec}}\approx 500

Model	LR	WD	Excluded	Self-consistent	First-order corrected	Dipole	$L_{1}$	Forces	HOMO–LUMO	ERIC
			(%)	energy (mHa)	energy (mHa)	(D)	(e)	(mHa/Å)	(mHa)	(ratio)
				$\|\hat{E}_{\mathrm{KS}}-E^{*}\|$	$\|\hat{E}-E^{*}\|$	$\\|\Delta\bm{\mu}\\|_{2}$	$\\|\Delta\rho\\|_{L^{1}}$	$\langle\\|\Delta\mathbf{F}_{a}\\|_{2}\rangle_{a}$	$\|\Delta\varepsilon_{\mathrm{gap}}\|$
Density-space models
$\mathbf{P}_{\mathrm{abs.}}$ (NON)	$3{\times}10^{-3}$	$10^{-2}$	100 [100, 100]	– . 000 0 (338) (245, 431)	– . 0000 00 (11.5) (8.0, 16)	– . 0000 00 (0.952) (0.76, 1.1)	– . 0000 0 (0.616) (0.55, 0.67)	– . 000 00 (10.1) (8.5, 12)	– . 000 00 (6.72) (5.8, 7.7)	– . 000 (0.830) (0.82, 0.84)
$\mathbf{P}_{\mathrm{abs.}}$ (NON), no rescale †	$3{\times}10^{-3}$	$10^{-2}$	100	– . 000 0 (175)	– . 0000 000 (4.07)	– . 0000 00 (0.594)	– . 0000 0 (0.477)	– . 000 000 (8.45)	– . 000 00 (5.36)	– . 000 (0.821)
$\mathbf{P}_{\mathrm{base}}$ (NON)	$3{\times}10^{-3}$	$10^{-2}$	100 [100, 100]	– . 000 00 (75.7) (61, 98)	– . 0000 000 (2.17) (1.3, 3.8)	– . 0000 00 (0.581) (0.50, 0.72)	– . 0000 0 (0.446) (0.41, 0.52)	– . 000 000 (7.51) (6.7, 9.0)	– . 000 00 (5.06) (4.9, 5.3)	– . 000 (0.816) (0.81, 0.82)

Appendix

Table 26: Observable errors under size extrapolation to QM40 for QM9-trained models ( table 5 ). Each cell reports the mean over seeds of the per-molecule mean absolute error, with the seed range in square brackets. Values in round brackets average all molecules of the lane before outlier exclusion, identical across models. Per column, the best bracketed mean and the best single seed are underlined. Predictions with DIIS residual

r_{\mathrm{DIIS}}/\sqrt{N_{\mathrm{elec}}}>30

mHa are excluded as outliers, where

r_{\mathrm{DIIS}}=\sqrt{8}\,\|\mathbf{S}^{-1/2}(\mathbf{F}\mathbf{P}\mathbf{S}-\mathbf{S}\mathbf{P}\mathbf{F})\mathbf{S}^{-1/2}\|_{F}

for scored density

\mathbf{P}

and Fock matrix

\mathbf{F}=\mathbf{F}[\mathbf{P}]

. Unbracketed values average kept molecules only, and Excluded reports the lost molecule percentage. Models losing all molecules display a dash before each bracketed value.

\hat{E}_{\mathrm{KS}}

is the KS-DFT energy evaluated at the predicted density,

\hat{E}

adds the first-order self-consistency correction ( section 5 ), and

E^{*}

is the converged reference energy. Before outlier exclusion, 1000 molecules enter Excluded,

\hat{E}_{\mathrm{KS}}

\hat{E}

, Dipole,

L_{1}

, and HOMO–LUMO, while 200 molecules enter Forces and ERIC. ERIC reports SCF iterations relative to the MINAO guess, averaged over the full lane with unconverged runs capped at 2. LR and WD denote peak learning rate and weight decay from the training configuration. † Single seed without range.

Model	LR	WD	Excluded	Self-consistent	Dipole	$L_{1}$	Forces	HOMO–LUMO	ERIC
			(%)	energy (mHa)	(D)	(e)	(mHa/Å)	(mHa)	(ratio)
				$\|\hat{E}_{\mathrm{KS}}-E^{*}\|$	$\\|\Delta\bm{\mu}\\|_{2}$	$\\|\Delta\rho\\|_{L^{1}}$	$\langle\\|\Delta\mathbf{F}_{a}\\|_{2}\rangle_{a}$	$\|\Delta\varepsilon_{\mathrm{gap}}\|$
Density-space models
$\mathbf{P}_{\mathrm{base}}$ (NON)	$3{\times}10^{-3}$	$10^{-2}$	100 [100, 100]	– . 00 (46920) (1214, 137230)	– . 000 (145) (8.3, 369)	– . 000 (12.9) (5.5, 25)	– . 00 0 (314) (69, 756)	– . 00 0 (76.3) (59, 104)	– . 000 (1.23) (1.1, 1.4)
$\mathbf{P}_{\mathrm{base}}$ (ON)	$3{\times}10^{-3}$	$3{\times}10^{-3}$	100 [100, 100]	– . 00 (13720) (887, 31790)	– . 000 0 (14.2) ( 1.0 , 31)	– . 000 0 (2.03) (1.3, 3.2)	– . 00 00 (60.1) (8.7, 154)	– . 00 0 (10.8) (8.5, 14)	– . 000 (0.817) (0.81, 0.83)
GROOT rot	$10^{-2}$	$10^{-2}$	0 61.5 [53, 66]	10.4 0 000 ( 36.9 ) [10, 11] ( 28 , 44)	0 0.279 00 (3.29) [0.27, 0.30] (3.3, 3.3)	0 0.329 0 ( 0.647 ) [0.32, 0.34] ( 0.62 , 0.69)	0 2.41 000 ( 3.27 ) [2.4, 2.4] ( 3.2 , 3.4)	0 2.17 00 ( 3.48 ) [2.1, 2.3] ( 3.5 , 3.5)	0 0.743 ( 0.776 ) [0.74, 0.74] ( 0.77 , 0.79)

Appendix

Table 27: Observable errors under size extrapolation to QMugs for QM9-trained models ( table 6 ). Each cell reports the mean over seeds of the per-molecule mean absolute error, with the seed range in square brackets. Values in round brackets average all molecules of the lane before outlier exclusion, identical across models. Per column, the best bracketed mean and the best single seed are underlined. Predictions with DIIS residual

r_{\mathrm{DIIS}}=\sqrt{8}\,\|\mathbf{S}^{-1/2}(\mathbf{F}\mathbf{P}\mathbf{S}-\mathbf{S}\mathbf{P}\mathbf{F})\mathbf{S}^{-1/2}\|_{F}

exceeding

30\sqrt{N_{\mathrm{elec}}}

\hat{E}_{\mathrm{KS}}

is the KS-DFT energy evaluated at the predicted density ( section 5 ), and

E^{*}

is the converged reference energy. Before outlier exclusion, 950 molecules enter Excluded,

\hat{E}_{\mathrm{KS}}

, Dipole,

L_{1}

, and HOMO–LUMO, while 190 molecules enter Forces and ERIC. ERIC reports SCF iterations relative to the MINAO guess, averaged over the full lane with unconverged runs capped at 2. LR and WD denote peak learning rate and weight decay from the training configuration. † Single seed without range.

Figure 19: Density-readout models on the QMugs out-of-distribution (OOD) test set, resolved over electron count. QMugs-trained 24-epoch checkpoints evaluated under the dataset and protocol of table 28 , using 100 molecules per 20-electron bin over

N_{\mathrm{elec}}\in[160,660)

, with forces and ERIC evaluated on the 500-molecule force subset. The models evaluate the

\Delta\mathbf{P}_{\mathrm{MINAO}}

readout under

\mathcal{L}_{\mathrm{NON}}

and under ON-Loss, and the two manifold-restricted variants GROOT rot and GROOT proj ( section 4 ). Markers are the per-bin mean absolute error of each model’s single seed and every error panel uses a logarithmic vertical axis; table 28 reports the corresponding dataset-level means. ERIC is the effective relative iteration count ( Eberhard et al., 2026a ) , where values below one mean fewer SCF iterations than the reference initial guess.

Figure 20: Hamiltonian-readout models on the QMugs OOD test set, resolved over electron count. Same checkpoints, dataset, protocol, binning, and aggregation as fig. 19 for the

\Delta\mathbf{H}_{\mathrm{SAP}}

readouts, comprising the AO-basis readout with a training-set fitted constant gauge under

\mathcal{L}_{\mathrm{NON}}

and the two ROCKET stages of eqs. 10 and 12 . Table 28 reports the corresponding dataset-level means.

Model	LR	WD	Excluded	Self-consistent	First-order corrected	Dipole	$L_{1}$	Forces	HOMO–LUMO	ERIC
			(%)	energy (mHa)	energy (mHa)	(D)	(e)	(mHa/Å)	(mHa)	(ratio)
				$\|\hat{E}_{\mathrm{KS}}-E^{*}\|$	$\|\hat{E}-E^{*}\|$	$\\|\Delta\bm{\mu}\\|_{2}$	$\\|\Delta\rho\\|_{L^{1}}$	$\langle\\|\Delta\mathbf{F}_{a}\\|_{2}\rangle_{a}$	$\|\Delta\varepsilon_{\mathrm{gap}}\|$
Supervised, 24 epochs
$\mathbf{P}_{\mathrm{base}}$ (NON)	$3{\times}10^{-3}$	$10^{-2}$	${>}99.9$	0 4.01 (245)	0.0142 0 (59.0)	0.127 0 0 (3.92)	0.0343 (1.43)	– . 000 0 (16.2)	0.123 0 (9.80)	– . 000 (0.818)
$\mathbf{P}_{\mathrm{base}}$ (ON)	$3{\times}10^{-3}$	$3{\times}10^{-3}$	0 53.2	10.1 0 0 (31.4)	0.140 0 00 (0.365)	0.0638 0 ( 0.143 )	0.0495 (0.155)	0 0.696 00 (2.01)	0.406 0 (1.10)	0 0.628 (0.668)
GROOT rot	$10^{-3}$	$10^{-2}$	00 1.12	0 3.22 00 (9.29)	0.0712 0 (36.8)	0.114 0 0 (4.69)	0.101 0 (0.313)	0 0.678 00 (3.11)	0.481 0 (1.18)	0 0.643 (0.654)

Appendix

Table 28: Observable errors under size extrapolation to QMugs for QMugs-trained models ( table 6 ). Each cell reports the per-molecule mean absolute error of a single seed. Values in round brackets average all molecules of the lane before outlier exclusion, identical across models. Per column, the best bracketed value within each training group is underlined. Predictions with DIIS residual

r_{\mathrm{DIIS}}=\sqrt{8}\,\|\mathbf{S}^{-1/2}(\mathbf{F}\mathbf{P}\mathbf{S}-\mathbf{S}\mathbf{P}\mathbf{F})\mathbf{S}^{-1/2}\|_{F}

exceeding

30\sqrt{N_{\mathrm{elec}}}

mHa are excluded as outliers, judged per checkpoint and molecule. Unbracketed values average kept molecules only, and Excluded reports the lost molecule percentage.

\hat{E}_{\mathrm{KS}}

is the KS-DFT energy evaluated at the predicted density,

\hat{E}

adds the first-order self-consistency correction ( section 5 ), and

E^{*}

is the converged reference energy. Before outlier exclusion, 2500 molecules enter Excluded,

\hat{E}_{\mathrm{KS}}

\hat{E}

, Dipole,

L_{1}

, and HOMO–LUMO, while 500 molecules enter Forces and ERIC. ERIC reports SCF iterations relative to the MINAO guess, averaged over the full lane with unconverged runs capped at 2. LR and WD denote peak learning rate and weight decay from the training configuration.

Model	LR	WD	Excluded	First-order corrected	Self-consistent	Dipole	$L_{1}$	Forces	HOMO–LUMO
			(%)	energy (mHa)	energy (mHa)	(D)	(e)	(mHa/Å)	(mHa)
				$\|\hat{E}-E^{*}\|$	$\|\hat{E}_{\mathrm{KS}}-E^{*}\|$	$\\|\Delta\bm{\mu}\\|_{2}$	$\\|\Delta\rho\\|_{L^{1}}$	$\langle\\|\Delta\mathbf{F}_{a}\\|_{2}\rangle_{a}$	$\|\Delta\varepsilon_{\mathrm{gap}}\|$
$\beta=120^{\circ}$
$\mathbf{P}_{\mathrm{base}}$ (NON)	$3{\times}10^{-3}$	$10^{-2}$	100	– . 00000 (0.204)	– . 0000 (13.5)	– . 0000 0 (0.187)	– . 0000 0 (0.148)	– . 00 0 (4.12)	– . 000 0 (3.32)
$\mathbf{P}_{\mathrm{base}}$ (ON)	$3{\times}10^{-3}$	$3{\times}10^{-3}$	00 0.0426	0 0.00634 ( 0.00643 )	0 4.34 00 0 (4.35)	0 0.108 0 0 ( 0.108 )	0 0.0572 0 ( 0.0572 )	0 1.27 0 ( 1.27 )	0 0.995 0 (0.995)
$\mathbf{H}_{\mathrm{base}}+c_{t}\mathbf{S}$ (NON)	$3{\times}10^{-4}$	$10^{-3}$	00 0.213	0 0.104 00 (0.604)	0 0.157 0 0 (1.26)	0 0.235 0 0 (0.285)	0 0.183 0 0 (0.188)	0 4.35 0 (4.54)	0 1.19 0 0 (1.29)

Appendix

Table 29: Observable errors under conformational transfer across three 3BPA fixed-dihedral scans for QMugs-trained models ( table 7 ). Because 3BPA samples a single molecule across a dense conformational grid rather than distinct chemical formulas, this test isolates conformational generalization from compositional variation. Each scan constitutes a distinct block of 2347 frames (

\beta=120^{\circ}

) or 2350 frames (

\beta=150^{\circ},180^{\circ}

), evaluated separately without pooling. Each cell reports the per-molecule mean absolute error of a single seed. Values in round brackets average all frames before outlier exclusion, identical across models. Per column and block, the best bracketed value is underlined.

\hat{E}_{\mathrm{KS}}

is the KS-DFT energy evaluated at the predicted density,

\hat{E}

adds the first-order self-consistency correction ( section 5 ), and

E^{*}

is the converged reference energy. Energy columns subtract the per-frame signed error median separately for each model and block. Predictions with DIIS residual

r_{\mathrm{DIIS}}=\sqrt{8}\,\|\mathbf{S}^{-1/2}(\mathbf{F}\mathbf{P}\mathbf{S}-\mathbf{S}\mathbf{P}\mathbf{F})\mathbf{S}^{-1/2}\|_{F}

exceeding

30\sqrt{N_{\mathrm{elec}}}

mHa are excluded as outliers, judged per checkpoint and molecule. Unbracketed values average kept molecules only, and Excluded reports the lost molecule percentage. Models losing all molecules display a dash before each bracketed value. LR and WD denote peak learning rate and weight decay from the training configuration. Baseline potentials PaiNN, NequIP, and MACE train on QMugs energies alone and predict energy directly, with forces obtained by negative gradients.

Model	LR	Excluded	First-order corrected	Self-consistent	Dipole	$L_{1}$	Forces	HOMO–LUMO	ERIC
		(%)	energy (mHa)	energy (mHa)	(D)	(e)	(mHa/Å)	(mHa)	(ratio)
			$\|\hat{E}-E^{*}\|$	$\|\hat{E}_{\mathrm{KS}}-E^{*}\|$	$\\|\Delta\bm{\mu}\\|_{2}$	$\\|\Delta\rho\\|_{L^{1}}$	$\langle\\|\Delta\mathbf{F}_{a}\\|_{2}\rangle_{a}$	$\|\Delta\varepsilon_{\mathrm{gap}}\|$
Before self-consistency training (the QMugs-trained parents)
$\mathbf{P}_{\mathrm{base}}$ (NON)	$3{\times}10^{-3}$	0 95.9	0 0.0363 (33.4)	13.3 00 (142)	0 0.0680 (2.09)	0 0.0396 (0.667)	0 2.80 (47.0)	0 0.784 (17.8)	0 0.769 (0.980)
$\mathbf{P}_{\mathrm{base}}$ (ON)	$3{\times}10^{-3}$	0 84.7	0 0.0386 0 ( 3.58 )	44.1 00 (116)	0 0.0416 ( 0.805 )	0 0.0323 ( 0.269 )	0 2.64 ( 21.8 )	0 0.455 0 ( 7.99 )	0 0.741 ( 0.896 )
GROOT rot	$10^{-3}$	0 72.4	0 0.202 0 (63.7)	0 0.574 0 (33.7)	0 0.146 0 (4.14)	0 0.0713 (1.01)	0 4.37 (66.4)	0 1.49 0 (23.9)	0 0.794 (0.999)

Appendix

Table 30: Observable errors on the Transition1x test split across self-consistency training epochs ( table 8 ). These reactive geometries represent chemistry that the QMugs-trained parent checkpoints of table 28 never encountered during pre-training. The first block evaluates the 24-epoch parent checkpoints before self-consistency training. Subsequent blocks evaluate models warm-started from each parent after 2, 4, and 6 epochs of occupied-virtual gradient training under a cosine schedule. The final block evaluates models trained for 6 epochs with the residual of Zhang et al. (2024) . No run uses weight decay or experience replay. LR denotes the peak learning rate selected on the validation split by the self-consistency residual alone ( table 22 ). Each cell reports the per-molecule mean absolute error of a single seed. Values in round brackets average all 1715 frames before outlier exclusion, and the best bracketed value per column and block is underlined.

\hat{E}_{\mathrm{KS}}

is the KS-DFT energy at the predicted density,

\hat{E}

adds the first-order self-consistency correction ( section 5 ), and

E^{*}

is the converged reference energy. Predictions with DIIS residual

r_{\mathrm{DIIS}}=\sqrt{8}\,\|\mathbf{S}^{-1/2}(\mathbf{F}\mathbf{P}\mathbf{S}-\mathbf{S}\mathbf{P}\mathbf{F})\mathbf{S}^{-1/2}\|_{F}

exceeding

30\sqrt{N_{\mathrm{elec}}}

mHa are excluded as outliers, judged per checkpoint and molecule. Unbracketed values average kept molecules only, and Excluded reports the lost molecule percentage. Models losing all molecules display a dash before each unbracketed value. ERIC reports SCF iterations relative to the MINAO guess, averaged over all frames with unconverged runs capped at 2.

Figure 21: Wall time to a converged SCF from ML and MINAO starts on out-of-distribution QMugs molecules. (a) Time from the prepared molecule to a converged SCF against the electron count, for the MINAO baseline (dashed) and for SCFs started from the predictions of GROOT rot , GROOT proj and ROCKET (third stage). The timed region covers the initial guess, the model input, the forward pass and the SCF; only the per-molecule system setup is excluded. (b) Solid lines: mean cost of one SCF cycle, i.e. SCF time divided by cycle count, dominated by the Fock build and including DIIS and diagonalization. Dotted lines: time from the forward pass to the starting density, including the copy to the CPU and, for ROCKET, the density matrix after diagonalization of the predicted Hamiltonian; the GROOT models start from their predicted density directly. Points are single molecules; lines join per-bin medians in (a) and per-bin means in (b) over 20-electron bins, without fitted curves. The set holds three random QMugs molecules per 20-electron bin in

[160,660)

, 75 in total, all beyond the 150-electron training range, computed with the reference settings of table 2 . Each molecule runs on one NVIDIA H200. One discarded warm-up round absorbs compilation and the grid and density-fitting setup; the four pipelines then alternate over 10 timed rounds, and each molecule contributes its minimum. All runs converge, and the ML-started SCFs reach the baseline total energy to within

3\times 10^{-11}\,E_{h}

. Relative to MINAO, the median time to convergence is 0.70 for both GROOT models and 0.72 for ROCKET.

Explore similar work

May 22, 2026physics.chem-ph

Multitask learning with semiempirical orbital charges enables sample-efficient MLIPs

Machine learning interatomic potentials (MLIPs) require generating computationally expensive, large-scale training datasets to accurately simulate materials and molecules. Incorporating electronic structure information using multitask learning improves sample efficiency, however, training on full Hamiltonian matrices, which scale quadratically with the number of atoms, is intractable for large datasets. In this work, we show that multitask learning utilizing orbitally resolved semiempirical charges significantly improves sample efficiency and accuracy in MLIPs. To efficiently predict orbital charges, we implement a specialized equivariant model, reducing charge prediction error compared to an invariant baseline. By augmenting training with computationally inexpensive GFN1-xTB orbital charges, which scale linearly with the number of atoms, our model achieves a 46% reduction in energy mean absolute error and requires five times less data to match the performance of energy-only models. Furthermore, our approach outperforms models trained on expensive density functional theory (DFT) atomic charges, capturing orbitally resolved electronic complexity and forcing the network to learn a physically accurate latent space that spontaneously clusters metals by shared chemical properties. Because orbital charges are only required during training, this approach preserves inference efficiency, providing a scalable recipe for developing accurate, data-efficient foundation models for complex chemical systems.

Ihor Neporozhnii, Sjoerd Hoogland, Oleksandr Voznyy

Department of Physical and Environmental Sciences, University of Toronto, 1065 Military Trail, Toronto, M1C 1A4, Ontario, Canada. · The Alliance for AI-Accelerated Materials Discovery, 10 King’s College Road, Toronto, M5S 3G4, Ontario, Canada. · Department of Electrical and Computer Engineering, University of Toronto, 10 King’s College Road, Toronto, M5S 3G4, Ontario, Canada.

Interatomic Potentials Quantum Chemistry Multi--Task Learning Near-Earth Objects New Tasks

Ask PDF

May 20, 2026cs.LG

TriForces: Augmenting Atomistic GNNs for Transferable Representations

Machine learning interatomic potentials (MLIPs) achieve excellent accuracy when trained on large Density Functional Theory (DFT) data. To be useful in practice, they must often be adapted to target chemistries using small and expensive task-specific datasets. However, MLIPs transfer inconsistently across domains, with representations that often loose accessible composition and structure information. To address this, we present TriForces, a model-agnostic three-stream framework that separates composition and structure information, combined with self-supervised learning to preserve transferable representations. TriForces improves performance on MatBench and QM9 over baselines without needing DFT labels and enables efficient similar structure retrieval through its learned latent space. On OMat24, in limited-data training regime, TriForces reduces energy MAE by 57% at 20K samples only and improves force MAE across sample sizes. We release pretrained TriForces variants across multiple MLIP architectures with code at https://github.com/Ramlaoui/triforces.

Ali Ramlaoui, Alexandre Duval, Hannah Bull +4

Entalpic, Paris, France · Universit´e Paris-Saclay, CentraleSup´elec, Inria, Gif-sur-Yvette, France

Interatomic Potentials Density Functional Theory Graph Neural Networks War

Ask PDF

Jul 2, 2026cs.LG

Beyond Adam: SOAP and Muon for Faster, Label-Efficient Training of Machine Learning Interatomic Potentials

Machine learning interatomic potentials (MLIPs) have become a hallmark of AI for scientific simulation. While efforts on new architectures and datasets have led to increasingly accurate and general models, the choice of optimizer for training has largely remained unexplored, defaulting to Adam and its variants in the community. Here, we implement and systematically compare a class of recently proposed matrix-structured optimizers, including Muon, SOAP, and the hybrid SOAP-Muon, for training NequIP and Allegro MLIP models. We find that these optimizers can substantially outperform Adam in both convergence speed and final accuracy. SOAP and SOAP-Muon emerge as robust and consistently strong methods, while Muon only provides partial gains relative to Adam. The improvements are particularly pronounced under partial force supervision. Our results indicate that optimizer choice is an overlooked yet impactful design axis for MLIPs.

Gil Harari, Yoel Zimmermann, Ola Tangen Kulseng +4

John A. Paulson School of Engineering and Applied Sciences, Harvard University, Cambridge, MA, USA · Robert Bosch LLC Research and Technology Center, Watertown, MA, USA

Interatomic Potentials Adam Muon Label Propagation

Ask PDF