AdaptNTK: Adaptive Uncertainty Quantification and Active Learning for Neural Network Potentials

Prajwal Ananth; Shuwen Yue

Source-linked AI summary

AdaptNTK: Adaptive Uncertainty Quantification and Active Learning for Neural Network Potentials

Prajwal Ananth, Shuwen Yue

arXiv:2609.00488v1Aug 31, 2026cs.LGphysics.chem-ph

TL;DR

ML interatomic potentials need reliable, efficient uncertainty estimates because extrapolative configurations can produce inaccurate forces, while conventional neural-network estimates often struggle with sequential batch redundancy. AdaptNTK uses regularized empirical-NTK feature-space uncertainty with label-free recursive updates, achieving strong force-error ranking and active-learning performance at lower cost than an ensemble. Its demonstrated scope remains limited by gas-phase evaluation, indirect force-error features, and omission of broader diversity- or chemistry-informed acquisition baselines.

Problem
ML interatomic potentials can deteriorate on out-of-distribution configurations, while conventional neural-network uncertainty estimates generally cannot update batch scores without additional fitting or retraining.
Method
AdaptNTK represents configurations with predicted-energy parameter gradients and measures regularized Mahalanobis uncertainty in empirical NTK feature space, enabling label-free recursive updates.
Results
AdaptNTK achieves the highest mean force-error correlations and the lowest force errors in active-learning experiments on rMD17 and Transition-1X, with risk–coverage comparable to a three-member ensemble.
Takeaways & Limitations
AdaptNTK provides a practical single-model foundation for data-efficient active learning with sequential uncertainty updates and lower reported Transition-1X timing than the ensemble.
Takeaways & Limitations
The evaluation covers gas-phase molecules, uses energy-gradient features that relate indirectly to force errors, and omits dedicated diversity-based or chemistry-informed acquisition strategies.

Abstract

from arXiv · show

Machine learning interatomic potentials bridge the gap between quantum chemical precision and classical computational speed, enabling molecular dynamics simulations with first-principles accuracy. Their reliability is often improved through active learning, which iteratively expands the training set by identifying uncertain, out-of-distribution configurations. Existing uncertainty-quantification methods often involve a trade-off between computational cost and reliability, and generally cannot account for redundancy as an acquisition batch is assembled. Here, we introduce AdaptNTK, a single-model framework that measures uncertainty as a regularized Mahalanobis distance in empirical neural tangent kernel (NTK) feature space. With the NTK features fixed during acquisition, the uncertainty depends on the acquired configurations but not their reference labels. This allows the uncertainty to be updated recursively after each selection without retraining, reducing redundancy within an acquisition batch. On held-out rMD17 data, AdaptNTK achieves the highest mean correlations with force errors (Spearman 0.68, Pearson 0.71) and matches a three-member ensemble in error retention. In active learning experiments, AdaptNTK achieves the lowest force errors across rMD17 and Transition-1X, with particularly strong performance on transition-state configurations in Transition-1X. AdaptNTK provides a 2.6-fold speedup per Transition-1X cycle relative to the ensemble, providing efficient single-model uncertainty estimation with sequential updates for data-efficient active learning.

1 Introduction

ML interatomic potentials can fail when simulations encounter configurations outside their training data, making reliable uncertainty estimates important for active learning. AdaptNTK addresses this challenge with single-model NTK uncertainty that supports label-free sequential updates and strong benchmark performance.

Out-of-distribution configurations can cause inaccurate forces and incorrect molecular-dynamics trajectories.
Neural-network potentials generally lack the closed-form uncertainty estimates available from Gaussian-process potentials.Existing neural-network approaches commonly use ensembles or approximate Bayesian methods.
AdaptNTK updates uncertainty after each acquisition without reference labels or network retraining, reducing redundancy within a batch.The fixed NTK features make uncertainty depend on selected configurations rather than their labels.
AdaptNTK measures uncertainty as a regularized distance from training data in empirical-NTK feature space.It represents configurations through gradients of predicted energy with respect to model parameters.
AdaptNTK achieves the highest mean force-error correlations and the lowest active-learning force errors on rMD17 and Transition-1X.Its normalized risk–coverage performance is comparable to a three-member ensemble, while reported Transition-1X timing is less than half the ensemble cost.

2 Related Work

Prior work includes ensemble, Bayesian, dropout, evidential, and reduced-cost uncertainty methods, while related active-learning methods often combine gradient features with diversity criteria. AdaptNTK instead uses a direct closed-form regularized quadratic form in empirical NTK feature space.

Deep-neural-network uncertainty methods include ensembles, approximate Bayesian approaches, dropout, learned-feature distributions, and evidential regression.
Empirical-NTK uncertainty has been studied as a lower-cost alternative that can match or outperform deep ensembles in general regression.
The related NTK approach is post-hoc and sampling-based, whereas AdaptNTK directly evaluates a closed-form regularized quadratic form.
Batch active-learning methods commonly pair last-layer or gradient-feature representations with diversity criteria.Related approaches have also been developed for interatomic potentials and concurrent NTK-based selection.

3 Background

The background frames predictive variance as an uncertainty measure and introduces the empirical NTK through parameter-gradient features. For a finite trained network, treating the empirical NTK as a covariance yields a locally linearized, directly computable uncertainty.

Gaussian-process predictive variance provides a measure of uncertainty used for acquisition.
The neural tangent feature map is the parameter gradient ϕθ(x), whose inner product defines the empirical NTK.
AdaptNTK uses the empirical NTK of a finite trained network rather than assuming the infinite-width regime.The construction corresponds to a local linearization around learned parameters.

4 Methods

AdaptNTK interprets empirical-NTK predictive uncertainty geometrically as a regularized feature-space distance, approximates it with a Gaussian sketch, and updates it recursively during sequential acquisition. These choices preserve acquisition decisions while reducing computational cost and batch redundancy.

NTK uncertainty: Empirical-NTK predictive variance can be expressed in parameter space using the training feature covariance matrix.The empirical NTK feature matrix contains tangent representations of training configurations.
NTK uncertainty: In the noiseless limit, uncertainty equals the squared distance of a query feature from the training-feature subspace.Components orthogonal to that subspace contribute directly to predictive variance.
NTK uncertainty: Regularized NTK uncertainty is a Mahalanobis norm that emphasizes weakly represented feature-space directions and suppresses well-represented ones.The regularization parameter λ is tunable.
NTK uncertainty: Energy-gradient features encode structure-dependent sensitivity of the learned potential, although they relate to force errors indirectly.
Sketching: A Gaussian sketch projects high-dimensional NTK features into a lower-dimensional space while approximately preserving uncertainty geometry.The theorem bounds uniform preservation using a dimension controlled by the ridge effective dimension.
Sketching: Once sketch error is below the relevant acquisition-score gaps, sketched and full-space rules select the same ordered batch.This explains why increasing the sketch dimension beyond a sufficient value need not change selected configurations.
Sequential acquisition: Fixed top-B uncertainty scores can select redundant configurations because acquiring one point may reduce nearby points’ uncertainty.This issue is especially relevant for correlated trajectory-generated candidate pools.
Sequential acquisition: AdaptNTK updates uncertainty after each selection without retraining by exploiting rank-one covariance changes and the Sherman–Morrison identity.

5 Experiments

Experiments evaluate AdaptNTK for uncertainty quality, active-learning data efficiency, reactive configurations, and computational overhead. Across rMD17 and Transition-1X, it provides strong error discrimination, lower force errors, and efficient sequential acquisition.

Uncertainty quantification: AdaptNTK’s NTK-based uncertainty shows the strongest linear and rank association with prediction error on held-out rMD17 configurations.Its correlations are ρp = 0.706 and ρs = 0.683.
Uncertainty quantification: AdaptNTK is the only method ranking among the top two for both discrimination metrics and ENCE.Calibration can improve uncertainty scale without changing ranking, so these metrics capture distinct properties.
Ablations: Sequential AdaptNTK outperforms the non-adaptive top-B baseline throughout Transition-1X acquisition.The comparison isolates the contribution of updating uncertainty after each selection.
Comparison of data efficiency: Across four rMD17 molecules, AdaptNTK attains the lowest mean force RMSE as the labeled set grows from 50 to 1,000 configurations.All methods use the same fixed seeds, data split, architecture, and optimization schedule.
Reactive configurations in Transition-1X: On Transition-1X, AdaptNTK reaches approximately 1 kcal mol−1 Å−1 force RMSE at 1,000 configurations, versus roughly 4 for QBC.The same ordering holds for energy and force errors on the transition-state split.
Computational overhead: AdaptNTK completes the reported Transition-1X timing run in 2,423 s versus 6,225 s for QBC.This corresponds to a 2.57× reduction in total runtime while retaining sequential candidate rescoring.

6 Limitations

The experiments are limited to gas-phase molecules and use energy-gradient features that relate to force errors only indirectly. Force errors may also poorly reflect downstream simulation performance, and broader acquisition strategies were not evaluated.

The experiments evaluate uncertainty quality and active-learning performance for gas-phase molecules, leaving extension to periodic systems as future work.
Energy-gradient features relate to force errors only indirectly; force-parameter Jacobians would be more direct but substantially larger.
Force errors can be poor indicators of downstream simulation performance, and the relationship to simulation-observable uncertainty remains uncharacterized.
The active-learning comparisons omit dedicated diversity-based and chemistry-informed acquisition strategies.
The experiments establish NTK-based uncertainty as an effective active-learning signal, not superiority over chemically informed or post-hoc selection heuristics.

7 Conclusion

The paper presents NTK-based uncertainty as a post-hoc, single-model estimate that supports sequential active learning. It reports strong error correlations, low active-learning errors, and lower cost than a three-member ensemble.

NTK-based uncertainty converts a trained neural potential into a post-hoc pointwise estimate without a committee or learned uncertainty head.
AdaptNTK achieves the strongest error correlations across four rMD17 molecules and nearly matches the three-member ensemble in error retention.
AdaptNTK updates uncertainty within acquisition batches without reference labels or additional model fitting.
AdaptNTK achieves the lowest reported errors in the rMD17 and Transition-1X active-learning studies.
The top-B ablation supports sequential updates, while the sketch-dimension sweep shows no observed benefit beyond p = 128.

A Proof of the sketch-preservation result

The proof establishes sketch preservation by analyzing the regularized kernel geometry and controlling Gaussian-sketch covariance. It then transfers the resulting score bound to the ordered active-learning path.

The appendix proves Theorem 4.1 and Corollary 4.2.
The regularized kernel is defined as K = ZZ⊤ and M = K + λI, with normalized features H = M−1/2Z.
The sketched feature vectors have covariance Σ = HH⊤, whose eigenvalues lie in [0, 1].
Gaussian sample-covariance concentration in effective dimension provides a high-probability approximation, made simultaneous over candidates by setting η = δ/m and applying a union bound.
The proof translates the approximation into a score bound using a Schur complement and the Woodbury identity.
For active learning, a union bound over step–candidate pairs shows that the sketched procedure follows the full-feature ordered greedy path.

B Ablation studies

The ablations separately test sequential updates and Gaussian-sketch dimension on Transition-1X. They use controlled protocols to assess whether within-batch adaptation and larger sketches change acquisition or final force error.

The ablation studies isolate sequential-update contributions and test sensitivity to Gaussian-sketch dimension.
Figure 4 compares adaptive AdaptNTK with a non-adaptive top-B NTK variant, QBC, and Random under shared experimental conditions.
Table 3 reports final force RMSE after the same 20-checkpoint Transition-1X protocol, with every setting ending at N = 1,000.
The maximum difference among sketch-dimension point estimates is 0.12 kcal mol−1 Å−1, and the ordering is non-monotone.
The reported run provides no evidence that a larger sketch improves final force error beyond p = 128.

C.1 Datasets and splits

The experiments use fixed molecular datasets and acquisition protocols to compare uncertainty-guided active learning, including sequential AdaptNTK selection and non-sequential baselines.

C.1 Datasets and splits: Each rMD17 uncertainty-quality dataset uses 1,000 training, 100 validation, and disjoint 2,000-configuration evaluation and calibration subsets.The four molecules are aspirin, naphthalene, paracetamol, and salicylic acid.
C.1 Datasets and splits: Active-learning runs hold out 25,000 test configurations, initialize with 50 training configurations, and add 50 configurations per checkpoint through 1,000 labeled configurations.The 20 checkpoints include the initial fit.
C.1 Datasets and splits: Transition-1X evaluation uses a transition-state split excluded from pretraining to test acquisition performance in chemistry relevant to reactivity.Learning curves use the same fixed seeds and extend through 1,000 labeled configurations.
C.1 Datasets and splits: At each checkpoint, methods fit the potential and uncertainty estimator, then select the next 50 configurations from the remaining pool.AdaptNTK updates scores after each selection without using the selected configuration’s label, whereas baselines select one highest-scoring batch.
C.1 Datasets and splits: The top-B ablation selects the 50 largest initial scores without sequential updates, while the sketch ablation changes only p.The sketch ablation keeps the data, model, regularization, acquisition budget, and fixed seed unchanged.

C.3 Interatomic potential and training

The study evaluates scalar uncertainty scores from several methods under a common neural-network potential setup, using force errors, retention, calibration, and timing metrics.

C.3 Interatomic potential and training: All methods use a ScaleShiftMACE model with 724,240 parameters, trained for 500 epochs with AdamW at learning rate 10^-4 and loss 0.01 LE + 0.99 LF.The architecture uses a 6.0 Å cutoff and batch size 16.
C.4 Uncertainty methods: AdaptNTK maps each configuration to a scalar uncertainty, while its NTK score uses a regularized Mahalanobis form in 512-dimensional sketched energy-gradient features.The feature covariance is computed from the training set.
C.4 Uncertainty methods: The ensemble baseline trains three identically configured models with different random seeds and scores uncertainty using force-prediction variability.The score averages standard deviations across atoms and Cartesian directions.
C.4 Uncertainty methods: MCD, SWAG, and evidential methods estimate uncertainty from stochastic forward passes, posterior weight samples, or an evidential regression head, respectively.MCD uses dropout at inference, SWAG retains up to 20 snapshots and uses 30 draws, and evidential uncertainty is β/(ν(α −1)).
C.5 Error target and metrics: Evaluation compares uncertainty with per-configuration force RMSE, using Pearson and Spearman correlations, normalized AURC, and cross-fitted ENCE.Normalized AURC uses oracle and random baselines, while ENCE averages relative RMSE deviations across ten equal-count bins.
C.3 Interatomic potential and training: Timing measurements use one NVIDIA A100 GPU and include pretraining, uncertainty construction, and batch selection but exclude reference-label costs.Label costs are identical for a fixed acquisition count.