Tensor cluster expansion as an efficient surrogate for multi-ACE

Technical Summary: Tensor Cluster Expansion as an Efficient Surrogate for Multi-ACE

Problem Statement
Training machine learning interatomic potentials (MLIPs), particularly those based on the Atomic Cluster Expansion (ACE) formalism like Multi-ACE (MACE), is computationally expensive. While fine-tuning foundation models reduces the need for massive new datasets, the generation of first-principles data (e.g., Density Functional Theory calculations) remains a bottleneck, especially for complex, concentrated alloys where the chemical design space explodes combinatorially. A critical step in this workflow is the strategic selection of a training set that maximizes chemical diversity. Ideally, this selection should be based on the same feature vectors used as inputs to the MLIP to ensure the training set maps to the model's output diversity. However, computing high-fidelity MLIP feature vectors (like MACE descriptors) for large candidate sets is often untractable on modest hardware, creating a barrier to efficient screening of chemically complex systems such as transition metal dichalcogenide (TMD) monolayers.

Methodology
The authors propose a framework where a Tensor Cluster Expansion (TCE) serves as a computationally cheaper surrogate for MACE features during the training set selection process.

1. Generalized TCE Framework: The authors generalize the on-lattice TCE method to arbitrary correlation orders ($m$-body terms). They define topological tensors ($A^{[\ell]}$) based on fully connected graphs ($K_m$) representing $m$-body clusters. These tensors are contracted with occupation tensors ($X$) to count species-decorated clusters. The resulting feature vector ($t$) consists of normalized cluster counts for various topological labels $[\ell]$.
2. Theoretical Relationship: The paper establishes that TCE features span a subspace of MACE features. Specifically, on a fixed lattice, the edges between neighbors (captured by TCE) can be recovered from a finite set of root-centered angular configurations used in MACE. This suggests TCE features can pre-screen the chemical diversity captured by MACE.
3. Maximum-Entropy Subset Selection: To select training configurations, the authors employ a greedy maximum-entropy subset selection algorithm. This algorithm iteratively selects configurations that maximize the determinant of the covariance matrix of the feature vectors, thereby maximizing the "footprint" of chemical diversity.
4. Experimental Workflow:
   • Dataset Generation: 100 random high-entropy TMD monolayer configurations ($(W, Mo, Ta)(S, Se, Te)_2$) were generated.
   • Surrogate vs. Target: Two subsets ($X_{TCE}$ and $X_{MACE}$) were selected using the entropy maximization algorithm on TCE and MACE feature vectors, respectively.
   • Potential Fine-tuning: The MACE-MP-0 foundation model was fine-tuned using Low-Rank Adaptation (LoRA) on both subsets independently.
   • Validation: The resulting potentials were tested on a held-out dataset to predict energies, forces, and temperature-dependent vibrational properties (Vibrational Density of States, VDOS) via Molecular Dynamics (MD) and Monte Carlo (MC) simulations.

Key Contributions

• Theoretical Generalization: The work extends the TCE formalism to arbitrary $m$-body orders and provides a proof (outlined in the text) that TCE features span a rotationally invariant subspace of MACE features.
• Efficient Surrogate Screening: The authors demonstrate that TCE features can act as an efficient, low-cost surrogate for MACE features in the context of maximum-entropy subset selection.
• Hardware Efficiency: The study highlights that TCE feature computation is significantly faster (approx. 17x on a single CPU) than MACE feature computation, enabling diversity sampling on resource-constrained hardware without requiring GPU acceleration.
• Validation on TMDs: The approach was successfully applied to alloyed TMD monolayers, showing that potentials fine-tuned on TCE-selected subsets yield vibrational properties nearly identical to those fine-tuned on MACE-selected subsets.

Results

• Computational Cost: Computing TCE features took approximately 0.18 seconds per configuration on a single CPU, compared to 3.14 seconds for MACE features, representing a ~17x speedup.
• Subset Overlap: The overlap between the subsets selected by TCE ($X_{TCE}$) and MACE ($X_{MACE}$) significantly exceeded random chance, indicating that TCE effectively identifies the same chemically diverse configurations as the more expensive MACE descriptors.
• Potential Accuracy: Both the TCE-selected and MACE-selected potentials showed comparable accuracy in predicting DFT energies and forces, with residuals symmetric around zero and no signs of overfitting.
• Vibrational Properties: When used to predict VDOSs across various temperatures (100 K to 400 K) for different TMD subsystems, the two potentials produced nearly identical spectral profiles. Peak locations were invariant with respect to the subset selection method, with only marginal differences in peak broadening.
• Chemical Ordering: The models successfully captured temperature-dependent chemical ordering, such as Te clustering in W(S, Te)$_2$ at low temperatures, which persisted at higher temperatures.

Significance and Claims
The paper claims that TCE features provide a tractable, cost-effective pathway for the pre-screening of chemically rich configurations for MLIP training. By using TCE as a surrogate, researchers can perform maximum-entropy subset selection on modest hardware (e.g., single CPUs) without sacrificing the diversity required to fine-tune high-fidelity MACE potentials. This approach facilitates the efficient prediction of finite-temperature alloy thermodynamics and vibrational properties in concentrated alloys. The authors position this work as a step toward "theory-in-the-loop" autonomous labs, where computationally efficient surrogates can guide the selection of samples for expensive ab initio calculations or MLIP training, thereby accelerating the discovery of materials for applications like clean energy and optoelectronics. The authors remain modest regarding limitations, noting that TCE is restricted to rotationally invariant features (limiting direct application to magnetic systems or equivariant vector properties) and that handling defects like interstitials requires further lattice augmentation.