Extending Quantum Transport Simulations to Realistic Device Dimensions through Machine-Learned Electronic Structures

Machine-learned (ML) operator models can be trained to predict ab initio Hamiltonian matrices at significantly reduced computational cost compared to density-functional theory (DFT), extending electronic structure calculations to previously unfeasible scales and allowing for their integration with advanced technology computer aided design (TCAD) tools. Here, we introduce MALOQ, a model to train on and predict electronic structure matrices for complex materials consisting of thousands of atoms, described by large orbital basis sets, and covering a wide range of atomic elements. MALOQ extends a state-of-the-art, SO(2)-equivariant architecture with custom data-processing kernels to handle high-rank sparse Hamiltonian matrix data and a scalable edge-wise distribution to treat large atomic graphs. We demonstrate training (inference) on structures made of 3,000 (48,000) atoms, up to 384 (128) GPUs. For systems with 10k+ atoms, MALOQ produces Hamiltonian matrices 10,000x faster than standard DFT codes, at high spectral accuracy. These matrices can be input to a quantum transport (QT) solver, enabling the simulation of entirely-atomistic devices consisting of 21,600 atoms (>180k orbitals). Overall, the full MALOQ+QT pipeline outperforms the reference DFT+QT one by a factor >150x, while maintaining comparable accuracy.