Gradient-Guided Furthest Point Sampling for Robust Training Set Selection

TL;DR

The paper tackles the data-inefficiency and robustness challenges in training-set selection for molecular PES modeling. It introduces Gradient Guided Furthest Point Sampling (GGFPS), which combines FPS with gradient norms $g_j=\| extbf{F}_joldsymbol ight rbracket o$ and a gradient-exponent schedule $\{eta_k\}$ to produce scores $s_j=(g_j)^{\beta_k} d_j$ and initial sampling probabilities $p_j=g_j/\sum_i g_i$. Empirical results on the 2D Styblinski-Tang function and MD17 trajectories show that GGFPS can achieve the full-dataset MAE with about half the training points on ST and significantly reduce MAE and predictive-variance across MD17, particularly for high-force-norm configurations. These findings demonstrate that gradient-aware sampling avoids FPS undersampling of equilibrium geometries and yields more robust, data-efficient surrogate models, with potential for extrapolative learning and diverse-system data generation. Kernel Ridge Regression is used as the predictive model, with Gaussian kernels for ST and a local FCHL19-based kernel for MD17, enabling effective evaluation of the sampling strategies.

Abstract

Smart training set selections procedures enable the reduction of data needs and improves predictive robustness in machine learning problems relevant to chemistry. We introduce Gradient Guided Furthest Point Sampling (GGFPS), a simple extension of Furthest Point Sampling (FPS) that leverages molecular force norms to guide efficient sampling of configurational spaces of molecules. Numerical evidence is presented for a toy-system (Styblinski-Tang function) as well as for molecular dynamics trajectories from the MD17 dataset. Compared to FPS and uniform sampling, our numerical results indicate superior data efficiency and robustness when using GGFPS. Distribution analysis of the MD17 data suggests that FPS systematically under-samples equilibrium geometries, resulting in large test errors for relaxed structures. GGFPS cures this artifact and (i) enables up to two fold reductions in training cost without sacrificing predictive accuracy compared to FPS in the 2-dimensional Styblinksi-Tang system, (ii) systematically lowers prediction errors for equilibrium as well as strained structures in MD17, and (iii) systematically decreases prediction error variances across all of the MD17 configuration spaces. These results suggest that gradient-aware sampling methods hold great promise as effective training set selection tools, and that naive use of FPS may result in imbalanced training and inconsistent prediction outcomes.

Figures (18)

• Figure 1: Left: A subset (orange) of labeled data (grey) from the Lennard-Jones potential, which represent a well performing training set from the labeled reaction data. Center: Furthest point (red), Boltzmann (dark blue), and gradient norm (yellow) sampling metrics applied to the Lennard-Jones potential. Right: Gradient norm sampling (yellow), and an instance of the GGFPS sampling metric (green).
• Figure 2: Left: A contour plot of the ST function surface in two dimensions. Center-left: The ST function gradient norm surface. Center-right: a heatmap of 50 ST function training sets generated by GGFPS, each with 100 data points, with a $\beta$ value of 0.5. In each of the three, a lighter color denotes a higher value. Right: ST function gradient norm distributions of training sets selected via URS (black dashed line), FPS (red dash-dotted line), GGFPS over increasing $\beta$ values (solid color-coded lines), and the constant $\beta^{\prime}$ versions of GGFPS (dotted lines). Training sets were sampled from 50 uniform randomly labeled sets of 1000 data points each.
• Figure 3: The MAE predictions of the 2D ST function surface with respect to training set size for URS, FPS, and GGFPS (black dashed, red dash-dotted and blue solid lines). The GGFPS $\beta$ values are optimized per training set fold via grid search. The FPS and GGFPS learning curves are read backwards. E.g. from a labeled set size of 1,000 data points (whose error is shown via URS), FPS and GGFPS sub-select from 950 to 50 training points. Note, at $N = 950$ the GGFPS learning curve is lower than the URS learning curve at $N = 1,000$. This is an artifact of using the MAE metric, and disappears when RMSE is used instead (See SI Figure \ref{['fig:SISTcombined']}).
• Figure 4: Box plots of the $\beta$ values corresponding to the lowest GGFPS CV errors with respect to training set size. Blue and orange box plots correspond to training sets drawn from labeled sets with 500 and 1,000 data points, respectively.
• Figure 5: Scatter plots of the 2D ST function absolute test errors corresponding to models trained on FPS (top row) and GGFPS (bottom row), with training set sizes, $N = 50$ (left column) and $N = 100$ (right column). All models start with an initial labeled set of 1,000 points. Median absolute test error values are shown above each scatter plot.