Physics-Inspired Augmentation of Equivariant Graph Neural Networks for Modeling Quantum Confinement

Modern nanoelectronic devices are complex heterostructures characterized by the central role that atomic-scale interactions and quantum mechanics play in their electronic properties. Simulating these devices using ab-initio methods, however, often proves to be intractable due to the cubic complexity of diagonalization of the quantum Hamiltonian. Machine learning methods including graph-neural networks (GNNs) have shown success overcoming this limitation in material systems dominated by local interactions, reproducing system properties with near ab-initio accuracy. In this work, we propose physics-inspired augmentations to GNNs to address nonlocal confinement effects -- a task that state-of-the-art GNNs fail at -- while maintaining favorable scaling in system size. We focus on nanosheets similar to those found in modern gate stacks, adding features to transmit boundary conditions through the structure. Our model demonstrates improved electronic density of states and current density of states predictions over the state-of-the-art and generalizes to structures beyond the training set, facilitating accurate, scalable predictions for a wide range of transistor designs.