Efficient tensorized evaluation of permutation invariant polynomials for representing potential energy surfaces
Abstract
Permutation invariant polynomials (PIPs), together with related polynomial-based invariant descriptors such as fundamental invariants (FIs), are widely used in constructing high-fidelity global potential energy surfaces (PESs) of molecules that contain identical atoms. Although the monomial symmetrization approach (MSA) enables fast evaluation of PIPs through recursive factorization, it leads to deeply nested computational graphs that may require large memory and are inefficient in modern automatic differentiation frameworks. In this work, we introduce JaxPIP, a JAX-based implementation that reformulates PIP/FI evaluation into tensorized linear algebra operations. By replacing recursive factorization with dense matrix operations combined with log-exp transformation and segmented summation, the evaluation becomes regular and GPU-friendly. This allows efficient execution with just-in-time compilation and enables large-scale batch evaluation of energies and forces (as well as higher-order derivatives). The resulting architecture supports ensemble simulations such as quasi-classical trajectory (QCT), path-integral molecular dynamics (PIMD), and diffusion Monte Carlo (DMC) in a fully vectorized manner. As demonstrated in examples, JaxPIP provides a practical route for efficient simulations of molecular systems with scalable GPU execution.
Related
- Representing chemical and enzymatic reactions in fingerprint space for applicability filtering and classification
- Permutation-inversion insensitive comparison of three-dimensional molecular structures
- Real-Time CASSCF Ehrenfest (TD-CASSCF) Modeling of Spintronics in a Tetrathiafulvalene-Based Nitronyl Nitroxide (BTBN) Model: The Competition between Charge Transfer and Spin Permutation.
- Amorphous Cu-FeOOH Clusters on Oxygen-Deficient Fe-Doped BiOCl Nanoflowers for Efficient Photo-Fenton Degradation of Levofloxacin
- Highly twinned cobalt selenide with abundant selenium vacancies for efficient nitrogen reduction reaction
BibTeX
@misc{li2026efficient,
title = {{Efficient tensorized evaluation of permutation invariant polynomials for representing potential energy surfaces}},
author = {Junhong Li and Kaisheng Song and Hua Guo and Jun Li},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.19584},
eprint = {2609.19584},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.19584},
note = {CC BY-NC-ND 4.0},
}
RIS
TY - UNPB
TI - Efficient tensorized evaluation of permutation invariant polynomials for representing potential energy surfaces
AU - Junhong Li
AU - Kaisheng Song
AU - Hua Guo
AU - Jun Li
PY - 2026
DA - 2026/09/17
PB - arXiv
DO - 10.48550/arxiv.2609.19584
UR - https://arxiv.org/abs/2609.19584
AB - Permutation invariant polynomials (PIPs), together with related polynomial-based invariant descriptors such as fundamental invariants (FIs), are widely used in constructing high-fidelity global potential energy surfaces (PESs) of molecules that contain identical atoms. Although the monomial symmetrization approach (MSA) enables fast evaluation of PIPs through recursive factorization, it leads to deeply nested computational graphs that may require large memory and are inefficient in modern automatic differentiation frameworks. In this work, we introduce JaxPIP, a JAX-based implementation that reformulates PIP/FI evaluation into tensorized linear algebra operations. By replacing recursive factorization with dense matrix operations combined with log-exp transformation and segmented summation, the evaluation becomes regular and GPU-friendly. This allows efficient execution with just-in-time compilation and enables large-scale batch evaluation of energies and forces (as well as higher-order derivatives). The resulting architecture supports ensemble simulations such as quasi-classical trajectory (QCT), path-integral molecular dynamics (PIMD), and diffusion Monte Carlo (DMC) in a fully vectorized manner. As demonstrated in examples, JaxPIP provides a practical route for efficient simulations of molecular systems with scalable GPU execution.
ER -