Measuring trainable degrees of freedom in materials graph neural networks: a random-subspace intrinsic dimension analysis
Authors: Shehroz Ahmad Shoaib, Kangming Li
Organizations: Physical Science and Engineering Division, King Abdullah University of Science and Technology, Thuwal, Saudi Arabia. · Department of Electrical Engineering, King Fahd University of Petroleum and Minerals, Dhahran, Saudi Arabia.
Final predictive accuracy is the standard basis for comparing graph neural networks (GNNs) in materials-property prediction, but it does not show how strongly performance depends on access to trainable parameter-space directions. Here, we introduce trainable-degree dependence as a complementary characterization of materials GNN learning. Using random-subspace intrinsic-dimension analysis, we train CGCNN, ALIGNN, and DimeNet++ in randomly oriented parameter subspaces across six prediction tasks and measure how performance recovers as independent trainable degrees of freedom are restored. The resulting recovery curves separate endpoint accuracy from the trainable-dimensional demand required to recover it. They reveal distinctions that final errors alone miss: metallic classification and log-bulk-modulus regression recover near-reference performance from small fractional subspaces, formation-energy and band-gap prediction show stronger architecture dependence, and phonon prediction is most sensitive to dimensional restriction. Dataset-size sweeps show that band-gap models require larger fractional subspaces as training data grows, whereas formation-energy and bulk-modulus responses are more stable. A width sweep shows that fractional thresholds can remain stable while absolute threshold dimensions increase with model size. Random-subspace analysis therefore provides a targeted stress test for how materials GNNs use their optimization space.
Figures & tables
Figure 1: Schematic of random-subspace training in a toy parameter space with ambient dimension D=3 . A model is initialized at θ0 , and optimization is restricted to an affine subspace θ0+Aϕ of dimension d≤D . The blue region represents parameter configurations that satisfy the target performance criterion, with θ∗ denoting a representative solution. As d increases, a randomly oriented subspace is more likely to intersect the acceptable-solution region; when d=D , the full parameter space is available for optimization.
Task
Data source
N
Dummy
ALIGNN
DimeNet++
CGCNN
Phonon frequency
matbench_phonons
1,265
323.8
73.4
87.3
108.0
Refractive index
matbench_dielectric
4,764
0.809
0.324
0.611
0.512
Bulk modulus
matbench_log_kvrh
110k
0.290
0.064
0.087
0.071
Band gap
MP23
153k
1.327
0.332
0.363
0.368
Formation energy
MP23
155k
1.059
0.048
0.060
0.053
Metallicity
matbench_is_metal
106k
0.500
0.949
0.948
0.937
Table 1: Dataset sizes and d=D predictive performance of the compact, approximately parameter-matched models. Regression tasks are evaluated using mean absolute error (MAE; lower is better), and the metallicity classification task using ROC-AUC (higher is better). The dummy reference is the dataset mean absolute deviation for regression tasks and 0.5 for metallicity classification.
Accurate prediction of crystal properties remains a key challenge in computational materials science. While graph neural networks (GNNs) such as CGCNN, MEGNet, ALIGNN, and SchNet have shown strong performance, they primarily represent crystals at the atomic level and implicitly learn local chemical environments through message passing. However, many material properties are governed by coordination polyhedra, the fundamental structural units formed by atoms and their neighboring atoms. To address this limitation, we propose the Coordination Polyhedron Graph Network (CPGN), a multi-scale GNN that jointly learns atomic, bond, and coordination-polyhedron representations. CPGN constructs three coupled graphs: an atom graph encoding elemental and bonding information, a line graph capturing angular interactions, and a coordination polyhedron graph describing Voronoi-derived local environments through corner-, edge-, and face-sharing relationships. Physically meaningful geometric descriptors are incorporated for each polyhedron, while an interleaved message-passing mechanism with bidirectional cross-attention enables effective information exchange across structural levels. Extensive evaluations on the Materials Project, JARVIS-DFT, and QM9 benchmark datasets demonstrate that CPGN outperforms existing state-of-the-art GNN models. It achieves a formation-energy MAE of 0.060 eV/atom and a band-gap MAE of 0.292 eV on the Materials Project, while providing competitive multi-property prediction on JARVIS-DFT and superior HOMO prediction on QM9. The results highlight that explicit modeling of coordination polyhedra improves crystal representation learning and enables accurate, physically interpretable prediction of material properties.
Sanjay Chakraborty
Department of Computer and Information Science (IDA), REAL, AIICS, Linköping University, Linköping, Sweden
Molecular message-passing neural networks commonly propagate chemically diverse interactions through a single graph, which may mix interaction-specific signals and require deep propagation to capture long-range effects. We introduce the Multi-level, Multi-color Graph Neural Network (MMGNN), a hierarchical framework that decomposes a molecular graph into overlapping atom-type-pair-specific subgraphs while preserving atom-level resolution. MMGNN-2D constructs chemical-colored subgraphs from covalent connectivity, whereas MMGNN-3D constructs geometric-colored subgraphs from spatial proximity and augments their edges with distance, angular, and torsional descriptors. Both variants apply a shared communicative message-passing backbone to each subgraph and combine the resulting representations through atom-wise aggregation and molecular readout. We evaluated MMGNN on five classification and three regression benchmarks from MoleculeNet using common scaffold splits and five independent runs. MMGNN-2D achieved the highest macro-average AUC-ROC of 0.838 across the classification datasets and the lowest RMSE on ESOL (0.803). MMGNN-3D obtained the highest mean AUC-ROC on BBBP (0.956) and the lowest RMSE on FreeSolv (1.793), indicating complementary strengths of topological and geometric representations. Structural and leave-one-out analyses further illustrate how the subgraph decomposition affects learned representations and atom-type-pair sensitivities. These results support overlapping interaction-specific graph decomposition as a competitive strategy for molecular property prediction.
Trung Nguyen, Duc Duy Nguyen
The Bredesen Center, University of Tennessee, Knoxville, TN 37996, USA · Department of Mathematics, University of Tennessee, Knoxville, TN 37996, USA
Graph neural networks (GNNs) have emerged as a versatile and efficient option for modeling the dynamic behavior of deformable materials. While GNNs generalize readily to arbitrary shapes, mesh topologies, and material parameters, existing architectures struggle to correctly predict the temporal evolution of key physical quantities such as linear and angular momentum. In this work, we propose MomentumGNN -- a novel architecture designed to accurately track momentum by construction. Unlike existing GNNs that output unconstrained nodal accelerations, our model predicts per-edge stretching and bending impulses which guarantee the preservation of linear and angular momentum. We train our network in an unsupervised fashion using a physics-based loss, and we show that our method outperforms baselines in a number of common scenarios where momentum plays a pivotal role.
Jiahong Wang, Logan Numerow, Stelian Coros +3
Max Planck Institute for Informatics · ETH Zurich · University of Bonn +1