Transplant Weights from Upstream Codes

Every state-of-the-art model from the literature, which is implemented in xnn, is verified against its reference code by weight transplant: MACE, NequIP, Allegro, CACE, PhysNet, ANI, BAMBOO, and the LES long-range add-on all load upstream weights and reproduce the upstream energies and forces to round-off error (see Model Fidelity to Upstream Codes for what each implementation matches, and the per-model notes below for the achieved precision). The one exception is SchNet: it is a clean-room build from the manuscripts, so there is no upstream code (or weights) to transplant; its verification reference is an independent NumPy implementation of the papers’ equations (examples/fidelity_checks/schnet_verification.ipynb).

The pattern

1. Build the xnn model with the same architecture hyperparameters as the upstream model (cutoff, l_max/max_ell, channels, layers, radial basis size, avg_num_neighbors, per-species energies/scales).

2. Map the upstream state dict onto the xnn parameter names.

3. load_state_dict and verify on a batch.

Transplants work in both directions: the same mapping loads xnn-trained weights back into the reference code (see the lock-step MD example below).

How much work step 2 is depends on the model:

• NequIP / Allegro: the parameter names already match (the interaction block uses upstream names), and Allegro’s strided channel-mixing linears keep the same flat weight layout as upstream, so a state dict trained with the reference code loads nearly as-is (and vice versa).

• MACE: the implementation reproduces upstream block by block, so whole models transplant through a simple name map. For the published pretrained checkpoints no manual transplant is needed at all: MACE.from_foundation("mace-mp-0-medium") (or an OFF23 alias, a URL, a local .model path, or an already-loaded mace-torch module) downloads, unpickles and converts the checkpoint in one call, including the ScaleShift energy expression, distance transforms, ZBL, the density interaction blocks, and multi-head slicing (head=...).

• CACE: same, with one gotcha: upstream lazily initializes some layers (the Bchi transform and the readout MLP), so run one forward pass on a sample batch before reading the upstream state dict.

• PhysNet: the transplant crosses frameworks (the original is TensorFlow 1.x); each TF variable is assigned from / to the corresponding torch parameter as a NumPy array.

• ANI: torchani’s pretrained ANI-1x, ANI-1ccx, and ANI-2x weights transplant into the ANI.ani1x() / ANI.ani1ccx() / ANI.ani2x() presets, for a single network or the full 8-model ensemble.

• BAMBOO / LES: plain name maps onto the xnn modules; for BAMBOO note the documented force-convention difference (xnn returns the conservative -dE/dr, equal to upstream forces + qeq_force).

Worked examples

The block-by-block notebooks in examples/fidelity_checks perform full transplants and check every intermediate tensor:

• mace_verification.ipynb: transplants a whole mace-torch model and reproduces its energy and forces to ~1e-15 (examples/gnn/mace/recreate_mace_architecture.ipynb is a longer companion tutorial that rebuilds the architecture step by step in both codes).

• nequip_verification.ipynb: ends with a whole-model weight transplant (~1e-16 agreement in energies, forces, and stress).

• allegro_verification.ipynb: the same for Allegro (~1e-15).

• cace_verification.ipynb: whole-model transplant from the original cace package, ~1e-16 relative in energies and forces, molecular and periodic.

• physnet_verification.ipynb: transplants the TF1 graph’s variables into the pure-PyTorch xnn model; energies, forces, and corrected charges match to ~1e-15 (float64).

• ani_verification.ipynb: transplants torchani’s pretrained ANI-1x / ANI-1ccx / ANI-2x ensembles; the AEV matches element for element to ~1e-16, ensemble energies to ~1e-8 Ha and forces to ~2e-7 Ha/Å.

• bamboo_verification.ipynb: transplants bytedance/bamboo weights; every GET layer, the charges, dipole, and component energies match to ~1e-15.

• les_verification.ipynb: transplants the upstream CACE-LR latent-charge head and Ewald settings into LatentEwald(CACE); each kernel matches to ~1e-16 in float64 and the whole model to float32 round-off.

A reverse transplant in production: the MD notebook examples/dnn/physnet/physnet_argon_density_md.ipynb loads xnn-trained PhysNet weights back into the original TF1 graph and propagates both engines through the same NVE trajectory in lock step.

The classical force fields need no transplant machinery at all: the force field is its parameter library, and the library format is the standard SEAMM .frc force-field file (Use and write .frc force-field files). ReaxFF loads the published fields shipped with xnn and ReaxFF-nn JSON libraries directly, and ReaxFF.export_library() writes trained parameters back out as .frc (classical) or JSON (with network weights). OPLS likewise loads the OPLS-AA distribution and its variants from .frc files – typing structures with the SMARTS templates those files carry – and OPLS.export_library() round-trips trained parameters through save_frc or the native JSON.

Tests

Parity with upstream given identical weights is also enforced in the test suite whenever the reference package is available: tests/test_nequip.py, tests/test_allegro.py, tests/test_cace.py, tests/test_ani.py (needs torchani, the [ani] extra), tests/test_les.py, tests/test_physnet.py (needs TensorFlow and an upstream clone via PHYSNET_UPSTREAM_PATH), and tests/test_bamboo.py (upstream clone via BAMBOO_UPSTREAM_PATH).