{"id":"d51aafd0-74ba-4b32-acab-dc1e1eb1b23a","arxiv_id":"2411.17011","paper_version":3,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":5.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":3,"one_line_summary":"MOLPIPx provides an end-to-end differentiable implementation of permutationally invariant polynomials for building machine-learned potential energy surfaces, with JAX and Rust backends.","lead":"MOLPIPx is a new open-source library for building machine learning models of molecular potential energy surfaces that automatically respect the symmetry of identical atoms. It computes energies and forces efficiently using two automatic differentiation engines, JAX for Python and EnzymeAD for Rust.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Automated MSA-to-JAX/Rust translation is the unverified load-bearing step: no test demonstrates permutational invariance of generated PIP vectors for all 15 symmetry types and degrees 3–7.","rationale":"The reader's weakest assumption is the same one I would flag: correctness of the automatically generated PIP code. I agree with it because the package's novelty is not the PIP algorithm—already established by MSA—but the conversion and AD integration. The paper does not supply a test that the conversion preserves symmetry; the examples show the package runs end-to-end, but not that the symmetry property holds. A wrong conversion is the kind of bug that would be silent: all models would fit the provided data (since each geometry is symmetrized consistently) but would fail on permuted inputs, invalidating molecular dynamics or any use requiring atom-order independence. The test I propose is cheap and would settle the question. I do not think this warrants rejection, because the implementation is open source and the underlying MSA algorithm is standard; it warrants a conditional acceptance requiring the verification. The Rust limitation and missing error bars are secondary concerns that also support conditional, but they are not the load-bearing issue.","tokens_in":21463,"tokens_out":7683,"duration_ms":71003,"concrete_test":"For every molecule in Table I and every degree p = 3,...,7, run the following verification using the released package: draw 100 random, non-degenerate Cartesian geometries; for each geometry generate all nontrivial permutations of identical atoms; compute Φ_PIP(x) and Φ_PIP(σx) with the JAX backend (and the Rust backend for the linear models it supports) and assert equality within 1e-8 relative tolerance; also assert that jax.jacobian(Φ_PIP) transforms correctly under the same permutations. Additionally, re-evaluate the original MSA .MONO/.POLY files on the same geometries and compare against the generated JAX/Rust outputs. If any mismatch appears, report the molecule and degree. This directly settles whether the automated conversion preserves permutational invariance.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central claim of the paper is that MOLPIPx reliably produces PIP vectors for 15 molecular symmetry types (Table I) and polynomial degrees 3–7, and that these vectors are permutationally invariant for each molecule. That property depends on one unverified step: the `msa_file_generator` translation of MSA's .MONO/.POLY files into JAX and Rust functions (Sec. III, Fig. 2). Any indexing or code-generation error in that translation would silently break permutational invariance for a specific molecule/degree, and every downstream linear, NN, or GP model built on that PIP vector would be wrong. The paper contains code listings only for A2B at degree 3 (Listings 2–3) and example results for CH4. It does not report any direct invariance test, such as checking equal energies (or equal Φ_PIP values, or equal forces) for permuted geometries, nor does it compare the JAX/Rust implementation against the reference MSA output. Because the translation is nontrivial—raw monomials are recombined by fpoly, with distance-order and mask conventions entering at several points—the risk is concrete rather than hypothetical. The paper also states in Sec. IV that the Rust backend is currently limited to linear models, so the abstract's 'higher-order derivatives' claim should be understood as JAX-only until demonstrated.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper presents MOLPIPx, an open-source software package that constructs permutationally invariant polynomial (PIP) basis vectors from MSA-generated monomial and polynomial files and exposes them through JAX and Rust/EnzymeAD automatic differentiation. The library is designed to support linear PIP models with energy and force fitting, PIP-based neural networks, Gaussian processes, and anisotropic Morse variables whose length scales are optimized by implicit differentiation of an inner least-squares problem. Demonstrations are carried out for methane using the external CH4 dataset of Ref. 28, and the paper reports force/energy prediction times, RMSE values, and optimization trajectories for lambda and hyperparameters.","tokens_in":127,"tokens_out":5970,"duration_ms":159108,"significance":"If the MSA-to-JAX/Rust translation is correct for all fifteen symmetry types and polynomial degrees listed, MOLPIPx fills a practical gap by removing the need to hand-code derivatives of PIP features and by making those features directly usable in standard ML frameworks such as Flax, GPJax, and Optax. The open-source release, the explicit code listings, and the use of external benchmark data are clear strengths; the benchmark data are independent of the package, so there is no evident circularity in the numerical results. The implicit-differentiation treatment of the Morse length scales is standard methodology rather than a new theoretical claim, but its integration with the PIP pipeline is useful. The main value of the paper is therefore software engineering: a flexible, differentiable PIP layer with automatic force computation.","major_comments":[{"comment":"The central load-bearing step is the automated translation of MSA .MONO/.POLY files into JAX and Rust functions, but the paper provides no direct verification of permutational invariance for the generated PIP vectors. Every downstream model in Eq. (1) inherits any indexing or code-generation error in fmono and fpoly; yet the only shown code is for A2B at degree 3, and the only tested molecule is CH4. I request a direct invariance test for all fifteen symmetry types in Table I and all degrees 3–7, for example by checking that Phi_PIP(x) is unchanged under permutations of like atoms, or by comparing the generated JAX/Rust output against the reference MSA output for a set of random geometries. In addition, Table I states that polynomial degrees range from 3 to 7, while Section IV.B reports timings for p = 8; the supported degree range should be stated consistently.","section":"Section III, Fig. 2, Table I, Listings 2–3"},{"comment":"The anisotropic Morse-variable formulation is presented as a contribution, but the paper does not report whether optimizing lambda_AA and lambda_AB improves accuracy relative to an isotropic lambda. Figure 4 shows optimization trajectories for methane, but it does not give final lambda values, does not compare the outer validation loss against an isotropic or fixed-lambda baseline, and does not report test-set RMSE for the optimized model. Without such a comparison, the claim that the anisotropic approach is beneficial rather than merely differentiable is unsupported.","section":"Sections II C and IV C, Eq. (11), Fig. 4"},{"comment":"The abstract states that MOLPIPx facilitates the computation of energy gradients and higher-order derivatives via JAX and EnzymeAD-Rust, but the only derivative computation demonstrated in the paper is first derivatives (forces). No Hessian or higher-order derivative example is provided for either backend, and Section IV explicitly states that the Rust version is currently limited to linear models. Please either demonstrate a higher-order derivative computation for at least one model (e.g., a Hessian calculation for a JAX linear or NN model) or qualify the abstract so that the validated scope is stated precisely.","section":"Abstract, Section IV.B, Section IV"}],"minor_comments":[{"comment":"The block matrix in Eq. (7) is hard to read and appears to contain indexing inconsistencies, for example rows labeled with partial derivatives of φ_np(x1) with respect to x3na but columns listing φ_np(x2). Please rewrite it with clear block notation that defines the geometry index and the ordering of the force rows.","section":"Section II B, Eq. (7)"},{"comment":"The constraint lambda in R is misleading because positivity is enforced later via softplus; please write the constraint as lambda in R^d (with positivity handled separately) or as lambda in R_{>0} once the softplus parameterization is introduced.","section":"Section II C, Eq. (11)"},{"comment":"The claim that the optimal value is lambda0 = 0.1 is based on a single grid-search run; please state the number of random seeds or repeated training runs used for each grid point, or otherwise explain why the comparison is robust to initialization noise.","section":"Section IV.D, Fig. 5"},{"comment":"The text says JAX jitting times for p = 8 are on the order of 10 hours, while the figure caption says the jitting time exceeded 5 hours for molecules with 5 atoms; please make the two statements consistent.","section":"Section IV.B, Fig. 3"},{"comment":"The decoration of a Flax Module class with @nn.jit may be confusing to readers unfamiliar with the exact Flax version used; a short note clarifying the intended Flax API version and the role of the decorator would help.","section":"Section III, Listing 4"},{"comment":"Timing and RMSE values are reported without uncertainties or replicate counts; either provide standard deviations over a small number of runs or explicitly state that the values are single-run measurements.","section":"Section IV, Figs. 3 and 7"},{"comment":"The entries of Table I are described as molecules, but most are abstract symmetry types such as A3 and ABCD; please clarify in the caption that these are generic atom-type labels rather than specific chemical species.","section":"Table I, Section III"},{"comment":"There are a few proofreading issues in the reference list, for example the text 'S. Manzhos, Sergei et al.' in the sentence preceding Ref. 91, and some entries mix author initials inconsistently; a careful proofread is recommended.","section":"References"}],"recommendation":"major_revision","confidential_remarks":"This is a software paper with a sound architecture and a useful open-source contribution. The main risk is the unverified correctness of the auto-generated symmetry code, but this is addressable with a focused verification section rather than by redesign. I recommend major revision; the requested invariance tests and anisotropic-versus-isotropic comparison are feasible within the scope of the manuscript."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Punchline: MOLPIPx is a real, useful package—it brings permutationally invariant polynomial models into JAX/Flax and Rust/EnzymeAD, with linear, NN, and GP wrappers and a small ARD-inspired anisotropic-Morse extension. The science is not new (PIP, PIP-NN, PIP-GP are prior art), but the implementation is a legitimate contribution: it makes end-to-end differentiable PIP fitting much easier.\n\nThe paper does a lot right. The code is open, the examples for CH4 are concrete, and the energy+force joint evaluation via AD is demonstrated. The authors are also honest that Rust currently only handles linear models, and the anisotropic λ optimization via implicit differentiation is a sensible use of prior work.\n\nSoft spots, in proportion: the stress-test concern is on target. The load-bearing step is the automated translation of MSA .MONO/.POLY files into JAX/Rust functions, and the paper shows no test that the generated PIP vectors are actually permutation invariant for all 15 symmetry types at degrees 3–7. The only listings are A2B at degree 3; the benchmarks use CH4. A code-generation bug for one molecule/degree would silently break every downstream model. This is fixable—a simple unit test comparing energies under atom permutations, plus a cross-check against MSA output—and the risk is medium, not high, because the example listing looks correct and MSA is well-established.\n\nSecond, the anisotropic λ is never compared to the isotropic baseline in terms of fitting accuracy, so we don't know if the extra parameters are worth using. The trajectories in Fig. 4 look fine, but they're not a benchmark. Third, the abstract's \"higher-order derivatives\" claim outruns the Rust backend, which is linear-only; the paper says this in Section IV, so it's a wording fix rather than a hidden flaw. The timing and RMSE numbers lacking uncertainties is minor for a software paper.\n\nWho it's for: people who build PIP-based PESs and want to switch to a differentiable, ML-forward workflow. I'd send it to peer review—it's exactly the kind of software contribution that deserves proper referee attention—and a good referee would ask for the invariance tests, the anisotropic baseline, and a more precise scope statement. I'd bring it to our reading group, and I'd cite it once the invariance tests are in (and even now as the tool description).","headline":"A genuinely useful JAX/Rust PIP package whose load-bearing MSA-to-code translation lacks a direct permutation-invariance test; send to review with that fixed.","tokens_in":22225,"tokens_out":2805,"would_cite":true,"duration_ms":26253,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"MOLPIPx makes permutationally invariant polynomial potential energy surfaces end-to-end differentiable, with automatic forces for linear, neural network, and Gaussian process models.","keywords":["permutationally invariant polynomials","potential energy surfaces","automatic differentiation","JAX","EnzymeAD","Gaussian process regression","neural network potentials","Morse variables"],"falsifier":"For each of the 15 symmetry types and each degree 3 to 7, evaluate the translated PIP vector on a random geometry and on the same geometry with identical-atom labels permuted; if any component of the vector or the resulting energy and force differs beyond numerical tolerance, the translation is not permutationally invariant and the central claim fails.","tokens_in":22,"feed_emoji":"⚛️","tokens_out":8488,"duration_ms":204120,"temperature":0.7,"pith_summary":"Potential energy surfaces describe how a molecule's energy changes with geometry, and permutationally invariant polynomials (PIPs) are a widely used way to build them because they respect the fact that swapping identical atoms leaves the energy unchanged. MOLPIPx is a software package that makes PIP-based models end-to-end differentiable: energies, forces, and higher derivatives are obtained automatically through automatic differentiation, removing the need for hand-coded gradients. It translates monomial and polynomial basis files, generated by the monomial symmetrization algorithm for 15 molecular symmetry types with polynomial degrees from 3 to 7, into Python (JAX) and Rust (EnzymeAD) code. The package also introduces anisotropic Morse variables with per-distance-type length scales optimized by gradient-based implicit differentiation. If this works as claimed, building and training accurate force fields for molecules becomes substantially more accessible.","feed_headline":"One library makes PIP energy surfaces fully differentiable","feed_subtitle":"MOLPIPx gives automatic energies and forces for 15 molecule types, linear models, neural nets, and Gaussian processes.","key_machinery":"The central object is the PIP vector $\\Phi_{\\mathrm{PIP}}(x) = (f_{\\mathrm{poly}} \\circ f_{\\mathrm{mono}} \\circ \\gamma \\circ d)(x)$. It maps Cartesian coordinates $x$ to inter-nuclear distances $r$, then to Morse variables $\\gamma = e^{-\\lambda r}$, then to symmetrized monomials via $f_{\\mathrm{mono}}$, and finally to symmetrized polynomials via $f_{\\mathrm{poly}}$; the symmetrization enforces invariance under permutation of like atoms. Automatic differentiation engines (JAX for Python, EnzymeAD for Rust) differentiate through this composition so that energies, forces, and higher derivatives are obtained as gradients. The package also introduces anisotropic Morse variables $\\gamma_{\\mathrm{aniso}} = \\exp(-\\sum_{\\ell} \\lambda_\\ell \\omega_\\ell \\odot r)$, where $\\omega_\\ell$ are binary mask vectors selecting each type of atom-atom distance, with $\\lambda_\\ell$ optimized by implicit differentiation.","core_discovery":"The authors claim that a PIP-based potential energy model can be expressed as a smooth composition of functions, $\\Phi_{\\mathrm{PIP}}(x) = (f_{\\mathrm{poly}} \\circ f_{\\mathrm{mono}} \\circ \\gamma \\circ d)(x)$, where $x$ are Cartesian coordinates, $d$ gives inter-nuclear distances, $\\gamma$ applies Morse variables $e^{-\\lambda r}$, and $f_{\\mathrm{mono}}$ and $f_{\\mathrm{poly}}$ build symmetrized monomials and polynomials. Because every layer is differentiable, the same $\\Phi_{\\mathrm{PIP}}$ vector feeds linear regression (with least-squares fitting that can include force Jacobians), feed-forward neural networks, and Gaussian processes with invariant kernels. Forces are obtained as gradients of the model output, and anisotropic Morse variables introduce per-type length scales whose outer/inner optimization uses implicit differentiation. The library ships pre-generated functions for 15 molecular symmetries and a converter that turns MSA-generated files into JAX or Rust code, with the Rust/EnzymeAD backend currently supporting linear models.","pith_inferences":["If the package is correct, PIP fitting can be embedded as a differentiable layer in larger end-to-end pipelines, for example training a force field jointly with downstream observables rather than fitting energies alone.","A natural extension is to treat the per-distance-type $\\lambda_\\ell$ as fully learned latent variables shared across a family of molecules, which the differentiable formulation makes straightforward.","The compile-time advantage of the Rust backend suggests ahead-of-time compilation could make very high-degree PIPs practical on larger molecules, a regime where JAX's jit cost is prohibitive.","A quick validation of permutational invariance on randomly permuted geometries for each of the 15 molecules would directly test whether the automated translation preserves the symmetry that the whole approach relies on."],"forward_implications":["A user can train linear PIP models, PIP neural networks, and PIP Gaussian processes from the same differentiable $\\Phi_{\\mathrm{PIP}}$ representation, with forces obtained as automatic gradients.","Force-augmented least squares for linear models becomes routine because the Jacobian $\\nabla_x \\Phi_{\\mathrm{PIP}}$ is available automatically.","The anisotropic Morse variables allow per-distance-type length scales to be learned by gradient descent with implicit differentiation, replacing grid search.","Post-compilation joint energy-force evaluation is faster than $10^{-3}$ seconds for low-degree models, and the Rust/EnzymeAD backend compiles large systems (degree 8, five atoms) in 33.5 seconds where JAX jit takes hours.","The same machinery transfers to many-body expansions, since the supplied symmetry types can be used term by term."],"supporting_citations":[{"why":"Introduces permutationally invariant polynomials in high dimensionality, the foundational representation MOLPIPx builds on.","marker":"[24]"},{"why":"Provides the standard review and framework of PIP-based potential energy surfaces that the library implements.","marker":"[25]"},{"why":"Supplies the CH4 training data and the force-augmented least-squares method used for linear PIP models.","marker":"[28]"},{"why":"Describes reverse-mode differentiation for PIP regression, the prior automatic-differentiation approach that MOLPIPx generalizes.","marker":"[30]"},{"why":"Defines the permutationally invariant polynomial neural network approach that MOLPIPx reproduces as PIP-NN.","marker":"[75]"},{"why":"Establishes the invariant Gaussian process kernel formulation that MOLPIPx implements as PIP-GP.","marker":"[95]"},{"why":"Provides the monomial symmetrization algorithm that generates the monomial and polynomial files the package translates.","marker":"[99]"},{"why":"Supplies the implicit differentiation method used for optimizing anisotropic Morse length scales.","marker":"[101]"},{"why":"Provides the JAX ecosystem that serves as the Python automatic differentiation backend.","marker":"[144]"},{"why":"Describes EnzymeAD, the compiler-based automatic differentiation engine used for the Rust backend.","marker":"[96]"}],"fun_headline_variants":["MOLPIPx: PIP energy surfaces go fully differentiable","Differentiable PIPs for energy surfaces in Python and Rust","MOLPIPx: one package for differentiable PIP energies and forces","PIPs meet JAX and Rust: differentiable energies from MOLPIPx","MOLPIPx makes PIP energy surfaces fully differentiable"],"cache_read_input_tokens":24448,"weakest_assumption_plain":"The load-bearing assumption is that the automated translation of MSA-generated monomial and polynomial files into JAX and Rust code produces correctly symmetrized PIP vectors for all 15 molecules and degrees 3 to 7, since no direct test of equal energies for permuted geometries is reported.","fun_headline_variants_meta":{"raw":{"variants":["MOLPIPx: PIP energy surfaces go fully differentiable","Differentiable PIPs for energy surfaces in Python and Rust","MOLPIPx: one package for differentiable PIP energies and forces","PIPs meet JAX and Rust: differentiable energies from MOLPIPx","MOLPIPx makes PIP energy surfaces fully differentiable"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000506,"raw_usage":{"total_tokens":2437,"prompt_tokens":880,"completion_tokens":1557,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":496,"completion_tokens_details":{"reasoning_tokens":1476}},"tokens_in":496,"tokens_out":1557,"duration_ms":10670,"temperature":1.0,"reasoning_tokens":1476,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-12T12:36:52.182413+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"For each of the 15 symmetry types and each degree 3 to 7, evaluate the translated PIP vector on a random geometry and on the same geometry with identical-atom labels permuted; if any component of the vector or the resulting energy and force differs beyond numerical tolerance, the translation is not permutationally invariant and the central claim fails.","supporting_citations":[{"cited_title":"Nandi , author C","cited_arxiv_id":null,"evidence_quote":"Supplies the CH4 training data and the force-augmented least-squares method used for linear PIP models."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Describes reverse-mode differentiation for PIP regression, the prior automatic-differentiation approach that MOLPIPx generalizes."},{"cited_title":"Qu , author Q","cited_arxiv_id":null,"evidence_quote":"Establishes the invariant Gaussian process kernel formulation that MOLPIPx implements as PIP-GP."},{"cited_title":"Xie \\ and\\ author J","cited_arxiv_id":null,"evidence_quote":"Provides the monomial symmetrization algorithm that generates the monomial and polynomial files the package translates."},{"cited_title":"Blondel , author Q","cited_arxiv_id":null,"evidence_quote":"Supplies the implicit differentiation method used for optimizing anisotropic Morse length scales."},{"cited_title":"Babuschkin , author K","cited_arxiv_id":null,"evidence_quote":"Provides the JAX ecosystem that serves as the Python automatic differentiation backend."},{"cited_title":"Moses \\ and\\ author V","cited_arxiv_id":null,"evidence_quote":"Describes EnzymeAD, the compiler-based automatic differentiation engine used for the Rust backend."}],"review_version":1}