LMHE-TSopt: Machine-Learned Leftmost Hessian Eigenvectors for Robust Transition State Finding

This repository contains the official implementation and associated scripts for the paper "Machine-Learned Leftmost Hessian Eigenvectors for Robust Transition State Finding". This codebase provides the framework for training $E(3)$-equivariant leftmost Hessian eigenvector (LMHE) predictors, integrating them into restricted-step partitioned rational function optimization (RS-PRFO), and executing robust transition state (TS) optimizations without the prohibitive cost of full Hessian calculations.

Repository Structure

1. Training (Train/)

Contains the implementation and training scripts for the $E(3)$-equivariant neural networks.

• GotenNet-GA/: Scripts for implementing and training the proposed architecture incorporating the global attention module.
• GotenNet/: Scripts for implementing and training the baseline (vanilla) GotenNet models.

2. Model Checkpoints (Models/)

Contains model parameters, hyperparameters, and training logs.

• GotenNet-GA/:
  • attnout1_logs/: Small-scale model checkpoints.
  • attnout3_logs/: Medium-scale model checkpoints (version_0 denotes an intermediate checkpoint; version_0-1 is the final converged model).
  • production_logs/: Ensemble of medium-scale models (version_n for $n=0$ to $4$ denote intermediate checkpoints; version_n-1 are the final ensemble members used for uncertainty quantification).
• GotenNet/:
  • vanilla_logs/: Small-scale baseline model.
  • vanilla1_logs/: Medium-scale baseline model.

3. Transition State Optimization (TSopt/)

Contains scripts for performing TS optimization with different strategies on the Sella benchmark set.

• qn/: Optimization using the standard quasi-Newton (TS-BFGS) updating method.
• full/: Optimization using exact analytical Hessians derived via automatic differentiation.
• ts-attnout3-version0-1/: LMHE-guided optimization using a single-inference predictor.
• ensemble5-ts-attnout3-version0-1/: LMHE-guided optimization incorporating the 5-model ensemble consistency check for real-time uncertainty quantification and dynamic fallback.

4. Analysis (Analysis/)

Jupyter notebooks for reproducing the quantitative evaluations presented in the manuscript.

• AnalyzeUncertain.ipynb: Evaluates ensemble consistency and its correlation with prediction accuracy.
• AnalyzeTS.ipynb: Classifies optimization results and analyzes wall times and gradient evaluation counts across methods.

Data Availability

The datasets used to train the LMHE predictors, along with the initial TS guess geometries and final optimized trajectories, are hosted on Figshare: 10.6084/m9.figshare.31791964.

Acknowledgements

The LMHE optimizer builds upon the innovative foundations provided by the following projects. Note that customized components extracted from GotenNet and e3nn are included directly within this repository.

• NewtonNet
• MLHessian-TSopt
• e3nn
• GotenNet
• Sella

Contact

For detailed information on setup and configuration, please feel free to open an issue in this repository or contact the corresponding authors.