Maximizing Efficiency of Dataset Compression for Machine Learning Potentials With Information Theory

TL;DR

This paper tackles the problem of optimizing training dataset size for machine-learning interatomic potentials without sacrificing accuracy or diversity. It introduces an information-theoretic framework that casts dataset compression as a minimum set cover (MSC) problem over atom-centered environments, and implements the QUESTS algorithm to greedily select a minimal yet information-rich subset of structures. Across GAP-20, TM23, and 64 ColabFit datasets, MSC consistently preserves outliers, maintains long-tail force distributions, and yields higher distribution overlap than baseline subsampling methods, with MLIPs trained on MSC-compressed data showing improved robustness in low-data regimes. The approach provides model-free, scalable metrics for compression efficiency and is publicly available in the QUESTS package, enabling subsampling, outlier detection, and efficient MLIP training at reduced cost.

Abstract

Machine learning interatomic potentials (MLIPs) balance high accuracy and lower costs compared to density functional theory calculations, but their performance often depends on the size and diversity of training datasets. Large datasets improve model accuracy and generalization but are computationally expensive to produce and train on, while smaller datasets risk discarding rare but important atomic environments and compromising MLIP accuracy/reliability. Here, we develop an information-theoretical framework to quantify the efficiency of dataset compression methods and propose an algorithm that maximizes this efficiency. By framing atomistic dataset compression as an instance of the minimum set cover (MSC) problem over atom-centered environments, our method identifies the smallest subset of structures that contains as much information as possible from the original dataset while pruning redundant information. The approach is extensively demonstrated on the GAP-20 and TM23 datasets, and validated on 64 varied datasets from the ColabFit repository. Across all cases, MSC consistently retains outliers, preserves dataset diversity, and reproduces the long-tail distributions of forces even at high compression rates, outperforming other subsampling methods. Furthermore, MLIPs trained on MSC-compressed datasets exhibit reduced error for out-of-distribution data even in low-data regimes. We explain these results using an outlier analysis and show that such quantitative conclusions could not be achieved with conventional dimensionality reduction methods. The algorithm is implemented in the open-source QUESTS package and can be used for several tasks in atomistic modeling, from data subsampling, outlier detection, and training improved MLIPs at a lower cost.

Figures (37)

• Figure 1: Figures of merit that quantify the efficiency of a dataset compression algorithm.a, An ideal algorithm should retain as much as the original dataset diversity as possible while removing redundancies from the data. In an entropy vs. diversity plot, this is depicted by curves that increase in entropy without reducing the diversity. At higher levels of compression, however, information loss is inevitable, and dataset diversity naturally start to decrease. b, An ideal algorithm should maximize the overlap of a compressed dataset with the original, complete dataset. Overlaps near 100% indicate a compression with effectively no information loss. c, An ideal dataset compression method should retain outlier environments as much as possible. Given that high-force environments tend to be outliers in datasets sampled from molecular dynamics or optimization trajectories, preserving this "long tail" of the force distribution is evidence of an algorithm that prioritizes these outlier environments. This long tail is represented by the cumulative distribution function (CDF) of forces.
• Figure 2: Performance of compression algorithms in subsampling three subsets of the GAP-20 dataset (Fullerenes, Nanotubes, Graphene). Each dataset was subsampled at four different sizes: 75, 50, 25, and 10%, and compared against the full, original dataset. a, our MSC method exhibits the best behavior in terms of simultaneous entropy increases and diversity retention for the selected datasets across dataset sizes. On the other hand, random sampling quickly leads to information loss, as shown by immediate diversity decreases upon subsampling. b, datasets compressed with MSC exhibit the higher overlap with the original dataset, with a single exception of Fullerenes sampled at 10% of its original size. c, MSC preserves the long tail of the distribution of forces of the environments, as shown by a higher cumulative distribution function (CDF) given a force threshold for datasets compressed at 25% of their original size (see also Fig. \ref{['fig:si:LongTailGAP20']}). d, Information entropy (orange) of the three subsets of GAP-20 and the maximum value of entropy that would be possible in a dataset with that size (gray). The numerical values of entropy (in nats) are shown with white numbers. The differences between the gray and orange bars show that the subsets exhibit different levels of redundancy, explaining the results in a--c. e, Average force error (eV/Å) above the minimum value of error given a subset (i.e., Graphene, Nanotubes, or Fullerenes) and a dataset size (i.e., 10, 25, 50, or 75%) across algorithms. While no algorithm consistently outperforms the others, models trained on datasets compressed using our MSC method exhibit a smaller range of errors above the minimum across sampling sizes, with errors more tightly grouped around optimal performance compared to other methods.
• Figure 3: Analysis of outlier loss across algorithms and dataset sizes for the three subsets of GAP-20 under study (Graphene, Nanotubes, and Fullerenes). a, the distribution of $\delta \mathcal{H}$ of environments in the full dataset w.r.t the compressed dataset across different methods indicates that MSC minimizes outlier loss, as shown by distributions of $\delta \mathcal{H}$ shifted towards less positive values. The only exception to this is the case of Fullerenes compressed at 10% of its original size, where the algorithm prioritized reducing the number of extreme outliers (i.e., those with $\delta \mathcal{H} > 10$ nats) at the expense of a reduction of true overlap ($\delta \mathcal{H} \leq 0$) with the dataset. b, a low-dimensional representation of the datasets, obtained with a UMAP visualization, illustrates how outlier loss is not trivially detected with these methods. When comparing the distribution of environments within the 10% and 75% datasets of Nanotubes, the information loss is at best qualitative, and at worst misleading. On all UMAP plots, the axes are the arbitrary axes of UMAP and brighter colors indicate a larger number of points within each region.
• Figure 4: Performance of subsampling algorithms across a subset of the TM23 dataset (elements Ir, Ag, Ti, Co). The datasets are subsampled 75, 50, 25, and 10% of their original sizes, and compared against the original datasets. a, our MSC method exhibits the best behavior in terms of simultaneous entropy increases and diversity retention for the subsets of TM23 across dataset sizes. In cases where there is lower data redundancy, such as the case of Ti or Co, information loss is inevitable. b, datasets compressed with MSC often exhibit the higher overlap with the original dataset compared to those obtained with other algorithms, except when redundancy in the original data is low, in which case all overlaps are similar. c, MSC preserves the long tail of the distribution of forces of the environments, as shown by a higher cumulative density function (CDF) given a force threshold $f$. d, Information entropy (orange) of the four subsets of TM23 and the maximum value of entropy that would be possible in a dataset with that size (gray). The numerical values of entropy are shown with white numbers, where the entropy is expressed in units of nats. The differences between the gray and orange bars show that the subsets contain different redundancy levels, explaining trends across elements in a--c. e, Average force error (eV/Å) above the minimum value of error given a subset (i.e., Ir, Ag, Ti, or Co) and a dataset size (i.e., 10, 25, 50, or 75%). Similarly to the GAP-20 example, while no algorithm consistently outperforms the other in terms of force errors, models trained on datasets compressed using our MSC method exhibit a smaller range of errors above the minimum across sampling sizes, with errors more tightly grouped around optimal performance in the low-data regime.
• Figure 5: Performance of baseline algorithms against MSC across 64 different atomistic datasets from the ColabFit repository compressed to 10, 25, 50, and 75% of their original sizes. Gray bars represent datasets for which the performance of MSC is worse than the baseline algorithm according to each metric. a, The difference in diversity between each method and MSC is mostly negative across all sizes and algorithms, demonstrating that MSC is a robust algorithm to preserve the original data distribution. b, Similarly, baseline algorithms lead to datasets with lower overlap against the original data compared to datasets subsampled with MSC. In the low-data regime, some algorithms showcase a higher overlap even though their diversity is lower. Similar to the Fullerenes example in Fig. \ref{['fig:umap']}, overlap does not fully account for how close outliers are to the distribution of points, but only for which points are entirely contained by a distribution.