MALOQ: Massively Accelerated Learning of Operators for Quantum Transport

2026-06-27

AI Summary

Traditional density functional theory struggles to efficiently simulate quantum transport systems comprising tens of thousands of atoms or more. This work proposes a machine learning framework based on SO(2)-equivariant neural networks, integrating high-rank Hamiltonian matrix processing kernels with edge-wise graph-based parallel distributed strategies to enable cross-element electronic structure modeling from small clusters up to systems with over 100,000 atoms. The method overcomes the efficiency bottleneck of molecular-scale distributed training, supporting scalable training and inference for systems of 3,000–12,000 atoms on the Alps supercomputer. It reduces per-epoch training time by over 30% compared to existing approaches and scales efficiently up to 256 GPUs, demonstrating the capability for high-throughput inference on material graphs of arbitrary size.

Abstract

Machine-learned (ML) operator models can be trained to predict density functional theory (DFT) Hamiltonian/density matrices at significantly reduced computational cost, thus extending electronic-structure calculations to previously unfeasible scales. Here, we introduce MALOQ (Massively Accelerated Learning of Operators for Quantum Transport), an application built to train on and predict electronic-structure matrices for systems made of few to 100k atoms, described by large basis sets, and covering a wide range of atomic elements. Based on a state-of-the-art, SO(2)-equivariant backbone architecture, MALOQ provides (i) custom data-processing kernels to handle high-rank Hamiltonian matrix data and (ii) a scalable edge-wise distribution of atomic graph(s). Trained on the largest molecular Hamiltonian datasets available today, it reduces time-per-epoch by over 30% compared to a molecule-wise-distributed framework, and enables inference on material graphs of arbitrary size. We demonstrate scalable training and inference for 3,000-12,000 atoms on the Alps supercomputer, up to 192 GPUs and 256 GPUs, respectively.

Problem

• quantum transport
• electronic-structure prediction
• Hamiltonian matrix
• large-scale systems
• density functional theory

Innovation

• operator learning
• SO(2)-equivariant architecture
• edge-wise graph distribution
• quantum transport
• scalable ML for electronic structure