PaperlayerView exact revision on arXiv

Source-linked AI summary

Hessian-based molecular conformation augmentation for a scalable and efficient strategy of machine learning interatomic potentials

Bumju Kwak, Jeonghee Jo

arXiv:2609.05233v1cs.LGphysics.chem-ph

TL;DR

MLIPs trained on energies and forces do not directly capture the PES Hessian needed for several applications, while explicit Hessian supervision adds computational and memory costs. This paper instead uses UniAug and ModeAug to generate Hessian-informed perturbations and Taylor-derived labels, reporting improved accuracy across equilibrium and non-equilibrium evaluations without changing the training objective or architecture.

  • Problem

    Energy-and-force-only MLIP training leaves Hessian information underused, while explicit Hessian supervision introduces higher-order differentiation and computational overhead.

  • Method

    UniAug and ModeAug use Hessian-derived isotropic or normal-mode perturbations with second-order Taylor expansions to synthesize labels without altering objectives or architectures.

  • Results

    Taylor-based augmentation enriches the PES and improves predictive accuracy across equilibrium, transition-state, and zero-shot molecular-dynamics evaluations.

  • Takeaways & Limitations

    The two schemes provide architecture-agnostic, plug-and-play alternatives to explicit Hessian supervision, with complementary advantages across training paradigms and target properties.

  • Takeaways & Limitations

    The distance cutoff used in the bond-length-deviation metric is a geometric heuristic.

Abstract

from arXiv · show

While machine-learning interatomic potentials (MLIPs) have successfully learned potential energy surfaces (PES) and atomic forces, many practical applications, such as vibrational analysis and transition state search, rely heavily on the PES Hessian. Yet, standard MLIPs tend to be trained on energy and forces alone, leaving Hessian information largely unexploited. Meanwhile, existing methods that explicitly incorporate the Hessian into training objectives require architectural modifications and introduce significant computational and memory overheads due to higher-order backpropagation. To address these limitations, we propose two Hessian-derived data augmentation schemes: isotropic Gaussian displacement (\textbf{UniAug}) and normal mode-weighted displacement (\textbf{ModeAug}). Both methods utilize simple Taylor expansions, achieving effective augmentation without altering training objectives or extending the autograd graph. This allows seamless, plug-and-play integration with existing architectures and training pipelines. Comprehensive evaluations across non-equilibrium and equilibrium datasets demonstrate that our approach enhances model accuracy while providing practical, task-specific guidelines.

Introduction

MLIPs learn energies and forces effectively, but applications requiring PES Hessians expose a gap that force-only training does not address. The paper proposes Hessian-derived Taylor data augmentation that improves models without changing their objectives or architectures.

  • Vibrational spectroscopy and reaction dynamics require accurate PES Hessians, which energy-and-force-only training fails to capture.
  • Explicit Hessian supervision incurs costly data generation, higher-order automatic differentiation, and substantial memory and computational overheads.
  • The proposed approach treats the Hessian as a data-generation prior, synthesizing perturbed geometries with energy and force labels from second-order Taylor expansions.
  • UniAug applies isotropic Cartesian displacements, whereas ModeAug samples Hessian-governed normal modes with amplitudes weighted inversely by eigenvalue magnitude.
  • Both schemes improve Hessian-free baselines across non-equilibrium and equilibrium benchmarks while leaving training objectives and model architectures unchanged.
  • UniAug and ModeAug provide complementary advantages depending on the training paradigm and target property.

Background

MLIP force modeling can be conservative or direct, while Hessian-based Taylor expansions provide a local route to synthesize energy and force labels. The approach depends on local smoothness and small displacements, and existing augmentation methods often leave labels expensive or unavailable.

  • Force prediction in MLIPs: Conservative force training obtains forces by differentiating predicted energy, guaranteeing a conservative field but requiring second-order automatic differentiation for force-loss backpropagation.
  • Force prediction in MLIPs: Direct force training can reduce training cost and memory, but its force Jacobian is generally asymmetric and does not define a valid energy Hessian.
  • Second-order Taylor approximation of the PES: The PES is smooth and the Taylor equations apply to arbitrary reference configurations, including off-equilibrium structures.
  • Second-order Taylor approximation of the PES: Around a reference configuration, second-order Taylor expansion uses the molecular Hessian and reference forces to approximate displaced energies and forces.
  • Second-order Taylor approximation of the PES: For small displacements, negligible higher-order anharmonic terms allow one reference Hessian to provide a locally consistent surrogate for energy and forces.
  • Data augmentation: Existing augmentation methods either produce inaccurate labels for perturbed structures or require separate quantum-chemical evaluations for every generated structure.

Methods

The method generates Hessian-informed perturbations with UniAug or ModeAug and assigns Taylor-derived energy and force labels. It evaluates these schemes across non-equilibrium, equilibrium, and molecular-dynamics settings using accuracy, structural, and rollout analyses.

  • Label synthesis: For both schemes, synthesized energy and force labels are computed from the second-order Taylor relation.
  • Augmentation schemes: UniAug applies isotropic Gaussian displacements to a randomly selected subset of atoms, with σ controlling displacement magnitude.
  • Augmentation schemes: ModeAug diagonalizes the mass-weighted Hessian, excludes translational and rotational modes, and weights vibrational amplitudes by inverse frequency.
  • Comparison of the two schemes: UniAug explores Cartesian space isotropically, whereas ModeAug gives larger displacements to low-frequency modes and smaller displacements to stiff modes.
  • Role of σ: Larger σ values broaden PES exploration but increase neglected higher-order Taylor contributions, so σ is selected separately by dataset and scheme.
  • Evaluation: The evaluation spans HORM, HessianQM9, and MDsim, including in-domain, zero-shot cross-domain, force-accuracy, structural, and rollout-stability analyses.
  • Training and evaluation: Models are trained in the E-F setting because augmentation changes only the data and not the loss function.

Results

UniAug and ModeAug alter sampled geometries and force shifts differently, while augmented training generally improves force, energy, Hessian, and eigenvalue prediction across in-domain, out-of-domain, equilibrium, and dynamical evaluations. Their relative performance depends on the augmentation scheme, perturbation scale, training paradigm, and task.

  • Augmentation behavior: 78% of UniAug samples at σ = 0.1 ˚A on ts1x-train exceeded ∆E > 0.1 eV, compared with 27% for ModeAug.The corresponding ModeAug fraction was 3% at the other reported scale.
  • Force-shift recovery: Augmented models generally recovered force-shift distributions more accurately than the No-aug model, while UniAug induced larger force variations than ModeAug.For equilibrium sets, both schemes produced distributions closer to ground truth and reduced zero-bias prediction collapse.
  • Augmentation behavior: UniAug produces broader, larger bond deformations—up to ∼15% versus within ∼5% for ModeAug—and correspondingly larger energy and force shifts.These differences are reported across all three HORM subsets at σ = 0.05 ˚A.
  • OOD evaluation: Both augmentation schemes reduced RGD1 force errors relative to E-F and No-aug baselines, although they did not match the original E-F-H force-error level.ModeAug was best under autograd force training across both energy and force metrics.
  • Equilibrium and dynamics: ModeAug achieved the lowest mean h(r) MAE, 0.0050 versus 0.0058 for UniAug and 0.0085 for No aug, while UniAug completed all 20 trajectories.ModeAug and No aug completed 16 and 14 trajectories, respectively; ModeAug also improved zero-shot equilibrium generalization.

Discussion

The discussion presents Taylor-expansion-based augmentation as a model-agnostic, low-overhead approach that improves coverage and predictive performance, while showing that the best scheme depends on the task and force-training paradigm. It also identifies practical selection criteria and scope limitations for local displacement regimes.

  • Method and efficiency: ModeAug and UniAug use Taylor-expansion-based data augmentation without modifying network architectures or loss formulations.The data-level intervention leaves the standard backpropagation graph unchanged and supports compatibility with existing MLIP pipelines.
  • Scope: The augmentation framework is particularly suited to off-equilibrium configurations where reference labels are sparse, but it is not designed for global conformational changes.Its intended regime is small displacement neighborhoods around reference configurations.
  • Data fidelity: Augmented energy and force shifts agree with reference DFT calculations, enriching the PES without introducing unphysical artifacts.The expanded coverage mitigates the equilibrium bias of standard relaxed datasets and translates into accuracy gains.
  • Comparative performance: Both augmentation schemes outperform the unaugmented baseline, but their relative effectiveness varies across benchmarks and evaluation targets.The optimal augmentation scheme depends on the force-training paradigm and the property being evaluated.
  • Comparative performance: UniAug is more advantageous for HORM in-domain force accuracy and energy accuracy, whereas ModeAug is stronger for HORM out-of-domain force prediction and generally slight Hessian improvements.For Hessian and eigenvalue reconstruction, ModeAug is slightly better in-domain while UniAug marginally leads out-of-domain; on HessianQM9, ModeAug performs slightly better overall.
  • Molecular dynamics: ModeAug better recovers interatomic pair-distance distributions in MD simulations, while UniAug provides superior trajectory completion stability.These results indicate that force prediction accuracy and Hessian reconstruction accuracy need not improve together.
  • Practical guidance: Choosing the augmentation mode and scale requires inspecting energy and force-shift distributions because broader, physically constrained coverage is associated with stronger model enhancement.The scale σ directly governs the distribution of synthesized configurations.
  • Limitations: The method is limited by heuristic pre-training selection, diminishing gains near zero-force equilibrium configurations, and the local validity of Taylor expansions.It is not intended to model global PES behavior across arbitrarily large conformational changes, including dihedral rotations or global vibrational modes.

Conclusion

The study introduces Hessian-derived UniAug and ModeAug, generating labeled configurations through second-order Taylor expansions without changing model architectures or objectives. Across HORM, HessianQM9, and MD17, augmentation improves predictive accuracy and reveals complementary strengths between the two schemes.

  • Method: UniAug and ModeAug generate energy and force labels for local configurations from a single Hessian without additional quantum chemical calculations.The approach injects curvature through synthesized labels while leaving the model architecture and objective function unchanged.
  • Results: Both augmentation schemes consistently enhance predictive accuracy across HORM, HessianQM9, and MD17.The benchmarks cover non-equilibrium reaction and transition-state structures, equilibrium configurations, and long-time dynamics.
  • Results: The approach recovers potential energy surface curvature that models trained only on energies and forces fail to capture.This directly targets curvature information that is otherwise absent from standard energy-force training.
  • Results: ModeAug is advantageous for curvature reconstruction and out-of-distribution generalization, whereas UniAug excels in energy accuracy and long-term simulation stability.The two augmentation modes therefore provide complementary benefits across evaluation metrics and tasks.
  • Implications: The study provides a computationally efficient route for leveraging curvature information in machine-learning potentials when Hessian data are available.The method avoids changing training objectives or adding training-time computational overhead.

Author Information

The authors state that they conducted the work independently of their employers and report no competing financial interest.

  • The authors conducted the work in their personal capacity, independently of and without support from their employers.
  • The authors declare no competing financial interest.

Supporting Information

The supporting information describes the HORM, HessianQM9, and MD17 benchmarks, their evaluation protocols, and the force-based training and simulation procedures. It also characterizes how non-equilibrium and equilibrium force distributions differ and explains the associated experimental scope.

  • HORM: HORM combines reaction pathways and transition states with DFT energies, forces, and full reference Hessian matrices.Its structures are non-equilibrium configurations with substantial forces, unlike databases focused mainly on relaxed geometries.
  • HORM: The HORM out-of-distribution evaluation uses RGD1 as a 60k test set for zero-shot generalization to reactions where up to two bonds form and break simultaneously.Models trained on ts1x-train are applied directly to RGD1 without fine-tuning, with augmentation strategy as the sole experimental variable.
  • Evaluation: The experiments use a joint energy-force MAE objective with identical architecture, seeds, and hyperparameters across augmentation strategies.Force target normalization is computed independently for each augmentation arm from its generated perturbation distribution.
  • HessianQM9: HessianQM9 force labels primarily reflect numerical convergence residuals rather than meaningful physical force signals.Their magnitudes are roughly two orders smaller than those in the off-equilibrium ts1x dataset, where forces are 0.36 eV/Å.

Implementation Details

Implementation uses EquiformerV2 with fixed training settings and applies UniAug or ModeAug through dynamically sampled perturbations. The same augmentation can also be performed offline without affecting model performance.

  • Model and training: EquiformerV2 is used as the primary SO(3)-equivariant backbone, with approximately 13.9M parameters and a 12 Å radial cutoff.The models use four layers, lmax = 4, mmax = 2, 128 spherical channels, and at most 20 neighbors.
  • Model and training: Training uses AdamW with weight decay 10^-3, gradient clipping at 10.0, and linear warmup followed by cosine learning-rate decay.All runs were conducted on a single NVIDIA GeForce RTX 3090 GPU with 24 GiB memory.
  • Augmentation settings: UniAug uses perturbation scales of σ = 0.05 Å for HORM and displaces fractions f = 0.125 on HORM and f = 0.25 on HessianQM9.For HessianQM9, σ is selected from a per-solvent range reported with the results.
  • Augmentation settings: ModeAug displaces all atoms through the normal-mode basis, while augmentation maintains a strict 1:1 balance between original and perturbed structures.Perturbations are sampled dynamically at the mini-batch level.
  • Data pipeline: Dynamic sampling is an implementation choice, because the augmentation strategies can also run offline during preprocessing without affecting model performance.
  • Experimental control: All baseline and augmented models use the identical joint energy-force loss, isolating data-level augmentation as the source of observed performance differences.

Bond Length Deviation Metric

The metric defines bonds once from covalent-radius distance criteria in the original geometry, then compares augmented structures using folded relative bond-length deviations. Its absolute chemical interpretation is limited near transition-state partial bonds, but paired UniAug–ModeAug comparisons remain robust.

  • Bond Definition: Bonds are identified once in the original geometry using element-pair distance cutoffs derived from covalent radii.The resulting cutoffs include 1.92 Å for C–C, 1.87 Å for C–N, 1.82 Å for C–O, 1.47 Å for C–H, 1.37 Å for O–H, and 1.02 Å for H–H.
  • Metric: Augmented structures retain the original bond list, and the metric reports folded relative deviations |∆r/r| for each bond.Using the same atomic pair indices enables paired comparisons between UniAug and ModeAug.
  • Validity and Limitations: The distance heuristic does not account for bond order, valence, or electronic structure.For the studied H, C, N, O, and F organic molecules, bonding and nonbonding distances are sufficiently separated for reliable bond identification.
  • Validity and Limitations: Transition-state partial bonds can lie near the 1.92 Å cutoff, so the bond definition does not align perfectly with chemical interpretation.Examples include forming C–C distances of 2.0 to 2.4 Å and C···H distances of 1.3 to 1.4 Å in hydrogen-transfer transition states.
  • Validity and Limitations: ModeAug produces a shorter tail in the |∆r/r| distribution, and this relative comparison is insensitive to the bond-definition details.The metric is therefore used for comparison between augmentation schemes rather than for absolute claims such as assigning a specific threshold to bond cleavage.

Effect of the Local Frame Detach on Autograd Force Evaluation

EquiformerV2 detaches its parameter-independent local rotation frame to reduce computational-graph cost, so autograd evaluates forces through a fixed-frame proxy rather than the full composite gradient. Removing the detach only at evaluation changes curvature error metrics appreciably.

  • Effect of the Local Frame Detach on Autograd Force Evaluation: The local rotation frame reduces SO(3) convolutions to SO(2) and is detached because it does not depend on model parameters.The detach reduces the cost of building the computational graph.
  • Effect of the Local Frame Detach on Autograd Force Evaluation: With the detach in place, autograd computes the gradient of E = f(x, R(x)) with the frame held fixed, rather than the true composite gradient.Training minimizes the force loss through this same fixed-frame path.
  • Effect of the Local Frame Detach on Autograd Force Evaluation: Removing the detach after training changes curvature error metrics appreciably at evaluation.The reported observation concerns evaluation behavior and is not used to argue for or against removing the detach.

Stiff Mode Frequency Calculation for Case Studies

The stiff-mode metric is computed from mass-weighted, symmetrized Hessians after removing translational and rotational modes. Signed vibrational frequencies are assigned to stiff or soft bands, and stiff-mode error averages frequency errors within the stiff band.

  • Stiff Mode Frequency Calculation for Case Studies: Predicted and reference Hessians are mass weighted as D = M^-1/2 H M^-1/2 using atomic masses repeated across Cartesian components.The mass-weighted Hessians are then diagonalized to obtain eigenvalues λ_i.
  • Stiff Mode Frequency Calculation for Case Studies: Because the predicted Hessian is not strictly symmetric, diagonalization uses its symmetric part 1/2(H + H^T).
  • Stiff Mode Frequency Calculation for Case Studies: The six eigenvalues with smallest |λ_i| are discarded as translational and rotational modes, leaving 3N − 6 vibrational modes.
  • Stiff Mode Frequency Calculation for Case Studies: Signed vibrational frequencies assign negative values to modes with λ_i < 0, representing imaginary frequencies.
  • Stiff Mode Frequency Calculation for Case Studies: Modes are classified using reference frequencies, with |ν_ref| < 400 cm^-1 defining the soft band; stiff-mode error is the mean absolute frequency error within the stiff band.
Loading 2609.05233v1…