Circumventing data imbalance in magnetic ground state data for magnetic moment predictions
Rohan Yuri Sanspeur and John R. Kitchin, “Circumventing data imbalance in magnetic ground state data for magnetic moment predictions”, Machine Learning: Science and Technology 5(1), 015023 (2024). https://doi.org/10.1088/2632-2153/ad23fb
This summary was written by a language model from the paper and has not been checked line by line. Rely on the published paper for details.
Question
Finding the magnetic ground state of a material with density functional theory (DFT) is expensive, both because each calculation is costly and because the number of possible magnetic orderings grows combinatorially. Machine learning surrogates could help, but magnetic training data are scarce and imbalanced. Earlier work treated the problem as classification of the collinear ordering of a whole structure, where most materials fall in the ferromagnetic and ferrimagnetic classes. This paper asks what happens when the problem is recast as regression of the magnetic moment on each atom, whether the imbalance survives that change, and how a graph neural network can be made to predict per-atom moments accurately anyway (Sanspeur 2024).
Approach
The authors queried the Materials Project for structures whose relaxed ground state is antiferromagnetic (AFM). This choice balances positive and negative moments and tests whether a graph neural network can resolve the sign of a moment. At the atom level, roughly 90% of the moments in this set sit near zero, giving a skewed continuous distribution with many zeros (Sanspeur 2024).
To handle this they borrow the two-part model idea from statistics on zero-inflated outcomes. In the first part, a GemNet-OC network is trained as a binary classifier that labels each atom magnetic or non-magnetic, using a cutoff of 0.1 µB on the moment magnitude, a class-weighted binary cross-entropy loss (scikit-learn’s balanced weighting), and 256-dimensional atom embeddings taken from the interaction blocks. In the second part, the embeddings of atoms classed as magnetic are passed to a stack of dense feed-forward blocks (linear layer, batch normalization, LeakyReLU, with dropout and residual connections) that regress the absolute value of the moment. The cutoff trades off removing the near-zero point mass against keeping enough samples, and its effect is examined in the supporting information on a larger set that also includes ferromagnetic (FM) and ferrimagnetic (FiM) structures (Sanspeur 2024).
Key findings
The classifier separates magnetic from non-magnetic atoms on the test set with an F1 score and Matthews correlation coefficient of about 91%, and a UMAP projection of the embeddings shows the two classes well separated while still varying within each class. A single regressor trained directly on all atoms is biased toward the mode and predicts roughly 0 µB, so its low error is misleading. The two-part scheme removes that bias and reaches a mean absolute error of 0.1 µB on moments from about 1 to about 7 µB (Sanspeur 2024).
The embeddings cannot tell spin-up from spin-down atoms. A ternary classifier that also predicted the sign of the moment produced overlapping embeddings for positive and negative moments, and regressing on the signed moment performed poorly. The authors attribute this to the symmetric local convolutions in the interaction blocks, which map equivalent environments to similar embeddings, and they sidestep it by regressing on the magnitude. When the classifier trained only on AFM structures was applied to FM and FiM structures (mostly metals and spinels, which are out of domain), the MCC fell to about 70%, but the embeddings still supported regression with errors similar to the AFM case (Sanspeur 2024).
As a case study, the models supplied initial moments (VASP’s MAGMOM) for the bulk oxide Mo2Ru13PtO32, where the oxidation-state and crystal-field logic of the WhereWulff workflow (Sanspeur et al. 2023) fails because Pt cannot be assigned an oxidation state. The classifier marked all atoms except Pt and O as magnetic. The regressor predicted 1.30 and 1.20 µB for the two Mo atoms and 0.63 to 1.08 µB for the Ru atoms. In single-point calculations, starting from these predictions reached the electronic ground state about 30% faster and gave an energy about 0.6 eV lower than starting from a uniform 5 µB (FM) or 0.6 µB (NM) initialization (Sanspeur 2024).
Significance
The paper shows that moving from structure-level ordering labels to atom-level moments does not remove data imbalance: it becomes a zero-inflated regression problem. A classify-then-regress two-part model is a simple fix that reuses the classifier’s learned embeddings. The work also points to a limitation of standard message-passing networks for magnetism: without an explicit spin feature they cannot represent collinear ordering. The authors propose adding the ordering to the initial atom embeddings and training moments and energies with a joint loss, and they suggest coupling the model to WhereWulff so that DFT is triggered only for out-of-distribution atoms (Sanspeur 2024). The practical result is better MAGMOM initial guesses for high-throughput DFT of magnetic oxides.