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

AI summary

Overview

• Research area: Computational materials science and machine learning, specifically dataset compression for machine learning interatomic potentials (MLIPs), combining Shannon information theory with combinatorial optimization.
• Technical level: Advanced (requires familiarity with information entropy, atom-centered descriptors, set cover problems, and neural network potential training).
• Scope: The paper proposes an information-theoretical framework for measuring how well dataset compression algorithms work, and introduces a minimum-set-cover algorithm that selects the smallest subset of atomic structures preserving the most information from the original dataset.

What This Paper Is About

Machine learning interatomic potentials can approximate density functional theory at lower cost, but their accuracy depends on the size and diversity of their training data. Large datasets are expensive to generate and train on, while aggressive subsampling risks throwing away rare but important atomic environments (outliers, metastable states, high-force configurations). This paper asks how to compress atomistic datasets so that redundancy is removed while as much information, diversity, and rare-event content as possible is retained.

Key Contributions

1. An information-theoretical framework for evaluating compression. The authors define figures of merit — information entropy, differential entropy, dataset diversity, dataset overlap, dataset efficiency, and the cumulative distribution function of force magnitudes — to quantify how much information a subsampling algorithm loses without needing to train any model.
2. A minimum set cover (MSC) compression algorithm. Atomistic dataset compression is recast as an instance of the minimum set cover problem over atom-centered environments, solved with a greedy polynomial-time approximation that iteratively picks the structure contributing the most uncovered environments while preferring diverse structures.
3. Extensive benchmarking against baselines. The method is compared against mean farthest point sampling (FPS), k-means clustering, and random selection on three GAP-20 subsets (Graphene, Nanotubes, Fullerenes), on the TM23 dataset, and validated on 64 varied datasets from the ColabFit repository.
4. A cautionary analysis of dimensionality reduction for compression diagnostics. Using UMAP visualizations, the authors show that 2D projections of subsampled datasets can be misleading about true information loss, and that per-environment differential entropy gives a quantitative answer instead.

Main Findings

• MSC wins on entropy and diversity: Across all metrics and datasets in the GAP-20 case study, the MSC method outperforms all baseline algorithms. It preserves dataset diversity as much as possible during compression and produces the highest increase in entropy, meaning redundancies are removed before information is lost. Random sampling shows immediate diversity loss upon subsampling.
• MSC maximizes overlap with the original dataset: MSC-compressed datasets have the highest overlap with the full dataset, with a single exception: the Fullerenes subset compressed to 10% of its original size, where all algorithms including MSC show large information loss.
• MSC preserves the long tail of force distributions: Compressed datasets retain a longer tail in the cumulative distribution function of force magnitudes, indicating broader coverage of the original dataset and better retention of outlier environments.
• Outlier retention is the mechanism: MSC shifts the distribution of differential entropy δℋ toward fewer positive values, meaning fewer environments in the original dataset become outliers relative to the compressed set. Fewer environments with δℋ > 0 means models are less likely to operate under out-of-distribution conditions.
• The 10% Fullerenes exception is explained: In that case MSC prioritized reducing extreme outliers (those with δℋ > 10 nats), leaving subsampled structures reasonably close to the original (0 < δℋ < 10 nats) but with lower strict overlap. This is interpreted as a consequence of Fullerenes having higher dataset efficiency than Nanotubes or Graphene, so higher compression implies information loss.
• Model training errors are more tightly grouped under MSC: No algorithm systematically achieves the lowest force error, but MLIPs trained on MSC-compressed datasets show a smaller range of errors above the minimum across sampling sizes. The effect is most pronounced at the 10% data regime, where MSC achieves the lowest possible errors among the three GAP-20 subsets.
• 2D visualizations are misleading: UMAP plots suggest random sampling covers more of the projected space, but the quantitative analysis shows random sampling has the worst information loss. The authors conclude that low-dimensional representations should not be used to judge dataset coverage.
• Broad validation: The abstract reports that across the GAP-20, TM23, and ColabFit cases, MSC consistently retains outliers, preserves diversity, and reproduces long-tail force distributions even at high compression rates, outperforming other subsampling methods.

Methodology in Plain English

Each atomic environment is described by an atom-centered descriptor that sorts distances from a central atom to its k = 32 nearest neighbors, plus a three-body term built from distances between those neighbors, using a cutoff of r_c = 5 Å. The amount of information in a set of these descriptors is measured with a kernel density estimate using a Gaussian kernel with bandwidth h = 0.015 Å⁻¹, calculated with the natural logarithm so units are nats. A related quantity, differential entropy δℋ, measures whether a new environment is already "contained" in a distribution (δℋ ≤ 0) or adds surprise (δℋ > 0).

The compression problem is framed as a set cover: the universe is all distinct atomic environments in the original data, and each structure is a set of environments. The goal is the smallest subfamily of structures whose environments span that universe. Since optimal set cover is NP-hard, a greedy approximation is used. The algorithm starts with the structure of highest per-structure information entropy, then repeatedly adds the structure maximizing the sum of the largest differential entropy among its environments plus its own entropy. This balances covering many uncovered environments against picking diverse rather than redundant structures.

For comparison, three baselines are implemented: random selection of structures, k-means clustering over per-structure averaged representations (using scikit-learn v. 1.7.0), and mean farthest point sampling, which repeatedly adds the structure whose mean vector is farthest from those already chosen.

To test whether compression hurts model accuracy, SevenNet models (v. 0.2.0, three equivariant blocks, L = 2, node features of 32 scalars, 32 vectors, and 32 tensors, cutoff 5 Å) were trained on datasets subsampled to 10%, 25%, 50%, 75%, and 100% of the original data. Each size was trained for a fixed number of structure-epochs: 10,000, 4,000, 2,000, 1,333, and 1,000 respectively. Batch size was 40, except at 10% where it was 4. Datasets were split 80:10:10 for train:validation:test with a single shared test set. Training used the Adam optimizer with initial learning rate 0.005 decaying by 0.99 per epoch and a Huber loss with force weight 0.1, stress weight 10⁻⁶, and δ = 0.01.

Why This Matters

Impact on research. Dataset construction is often the most expensive part of building a machine learning potential. This work provides model-free metrics that quantify information loss from compression without running expensive training sweeps, giving researchers an architecture-independent way to choose a compression level. It also shows that compression quality has its largest effect in low-data regimes, where the difference between algorithms translates into measurable generalization error.

Real-world applications (derived from the roles the paper assigns to MLIPs and atomistic modeling):

• Training interatomic potentials for materials screening and simulation at a fraction of the data-generation and training cost.
• Data subsampling of large simulation datasets prior to training, reducing memory footprint and compute.
• Outlier detection in atomistic datasets, surfacing rare events, metastable states, and high-force configurations that matter for reliability.
• Curating and validating datasets in community repositories, using overlap and diversity metrics as quality checks.

Industry relevance. Reducing the cost of producing and training on atomistic datasets directly affects organizations that generate DFT data at scale — battery and catalyst developers, semiconductor and alloy manufacturers, and any group running large molecular dynamics campaigns. The method is implemented in the open-source QUESTS package (https://github.com/dskoda/quests), and the trained models use the publicly available SevenNet code, so adoption does not require proprietary tooling.

Future Directions

• The paper notes that most baselines assume entire structures must be added to training sets; the authors state that all algorithms could instead operate per-atom rather than per-structure, which remains an avenue for further efficiency gains.
• The greedy approximation is used because optimal set cover is NP-hard; better approximations or exact solvers for restricted cases are not explored in the content provided.
• Because set cover operates in a continuous space rather than a discrete one, the mapping from the atomistic problem to the classical formulation is described as differing slightly from the original problem — a formal treatment of this gap is left open.
• The paper leaves open how the compression framework behaves when the cost of adding a data point (labeling cost) differs from the atom-count cost used here, and the TM23 results section is truncated in the provided content, so the full extent of that case study is not reported here.

Target Audience

This paper is aimed at machine learning researchers and computational materials scientists building or curating datasets for interatomic potentials, and at practitioners who train MLIPs on large simulation datasets and need principled subsampling tools. Readers with a background in information theory, atomistic simulation, or dataset engineering will get the most from it; the methodology section assumes comfort with kernel density estimation, entropy, and graph-like combinatorial formulations.

Authors’ 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.