[Paper Review] Connectivity Optimized Nested Graph Networks for Crystal Structures

Robin Ruff, Patrick Reiser+2 more|arXiv (Cornell University)|Feb 27, 2023

Machine Learning in Materials Science9 citations

TL;DR

The paper introduces asymmetric unit graph representations and a connectivity-optimized multi-graph network (coGN) for crystalline materials, then extends to nested graph networks (NGN) to systematically explore architecture space, achieving strong results on MatBench benchmarks.

ABSTRACT

Graph neural networks (GNNs) have been applied to a large variety of applications in materials science and chemistry. Here, we recapitulate the graph construction for crystalline (periodic) materials and investigate its impact on the GNNs model performance. We suggest the asymmetric unit cell as a representation to reduce the number of atoms by using all symmetries of the system. This substantially reduced the computational cost and thus time needed to train large graph neural networks without any loss in accuracy. Furthermore, with a simple but systematically built GNN architecture based on message passing and line graph templates, we introduce a general architecture (Nested Graph Network, NGN) that is applicable to a wide range of tasks. We show that our suggested models systematically improve state-of-the-art results across all tasks within the MatBench benchmark. Further analysis shows that optimized connectivity and deeper message functions are responsible for the improvement. Asymmetric unit cells and connectivity optimization can be generally applied to (crystal) graph networks, while our suggested nested graph framework will open new ways of systematic comparison of GNN architectures.

Motivation & Objective

• Reduce computational cost and redundancy in crystal GNNs by using asymmetric unit cells while preserving accuracy.
• Systematically explore graph neural network design space for crystals via Nested Graph Networks (NGN).
• Improve state-of-the-art on MatBench crystal-structure tasks through optimized connectivity and architecture choices.
• Provide practical guidelines on edge selection, symmetry exploitation, and line-graph based angle information for crystal GNNs.

Proposed method

• Represent crystals using asymmetric unit graphs to exploit symmetries and reduce node/edge counts.
• Compare edge selection strategies (k-NN, radius, Voronoi) for crystal graphs and assess impact on performance.
• Develop Nested Graph Networks (NGN) by integrating line-graph based angular information with graph networks, enabling multi-level message passing.
• Instantiate coGN (connectivity-optimized GN) with deep edge updates and connectivity optimization to maximize predictive accuracy.
• Conduct architecture search within the NGN framework, including variants that nest GN blocks on G, L(G), and L(L(G)).
• Evaluate on MatBench crystal-structure tasks with hyperparameter optimization on log_gvrh dataset and test transferability across tasks.

(a) $k$ NN Edges

Experimental results

Research questions

• RQ1 Does exploiting crystal symmetries via asymmetric unit graphs reduce computation without harming accuracy?
• RQ2 How does edge connectivity and choice of line-graph based angular features affect GNN performance on crystalline properties?
• RQ3 Can Nested Graph Networks subsume or improve upon existing crystal GNN architectures, and under what conditions?
• RQ4 What architecture and preprocessing settings yield state-of-the-art results on MatBench structure datasets?
• RQ5 Are the benefits of NGNs consistent across tasks or require task-specific hyperparameter tuning?

Key findings

• Asymmetric unit graphs reduce nodes/edges and memory footprint by about a factor of ~2.1 on MatBench datasets, with no loss in accuracy for E(3)-invariant GNNs.
• For k-NN edge selection with k=24, graphs achieve strong predictive performance; Voronoi-based edges also viable but with different connectivity.
• The coGN model with connectivity optimization achieves state-of-the-art results on multiple MatBench structure datasets, surpassing several prior models.
• NGNs offer deeper edge updates and line-graph based angle information, providing small but consistent improvements in some variants, yet plain GNs with optimized connectivity can match or outperform NGN configurations on several tasks.
• Optimal performance on log_gvrh was achieved with a DimeNet-like (Variant 2) architecture via hyperparameter optimization, indicating architecture choice and connectivity interactively determine gains.
• NGN training is more computationally intensive due to line-graph construction and larger parameter counts, motivating a trade-off between nesting depth and connectivity.

(b) Radius Edges