Enhancing the Scalability and Applicability of Kohn-Sham Hamiltonians for Molecular Systems

2025-02-26

AI Summary

To address the computational bottlenecks of Kohn–Sham density functional theory (DFT)—including high Hamiltonian construction cost, slow or non-convergent self-consistent field (SCF) iterations, and inaccurate ground-state property predictions—this work proposes a physics-informed deep learning framework. We introduce a novel wavefunction alignment loss (WALoss) that jointly optimizes Hamiltonian prediction accuracy and orbital energy errors, while explicitly enforcing rotational and translational symmetry to ensure physical consistency. Evaluated on the large-scale PubChemQH molecular dataset, our method reduces total energy prediction error by 1347× compared to conventional DFT approximations and accelerates SCF convergence by 18%. These improvements significantly enhance the scalability and physical fidelity of DFT for thousand-atom molecular systems.

Technology Category

• Machine Learning: Quantum Machine Learning
• Search and Optimization: Learning to Search
• Data Mining & Knowledge Management: Scalability, Parallel & Distributed Systems

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Density Functional Theory (DFT) is a pivotal method within quantum chemistry and materials science, with its core involving the construction and solution of the Kohn-Sham Hamiltonian. Despite its importance, the application of DFT is frequently limited by the substantial computational resources required to construct the Kohn-Sham Hamiltonian. In response to these limitations, current research has employed deep-learning models to efficiently predict molecular and solid Hamiltonians, with roto-translational symmetries encoded in their neural networks. However, the scalability of prior models may be problematic when applied to large molecules, resulting in non-physical predictions of ground-state properties. In this study, we generate a substantially larger training set (PubChemQH) than used previously and use it to create a scalable model for DFT calculations with physical accuracy. For our model, we introduce a loss function derived from physical principles, which we call Wavefunction Alignment Loss (WALoss). WALoss involves performing a basis change on the predicted Hamiltonian to align it with the observed one; thus, the resulting differences can serve as a surrogate for orbital energy differences, allowing models to make better predictions for molecular orbitals and total energies than previously possible. WALoss also substantially accelerates self-consistent-field (SCF) DFT calculations. Here, we show it achieves a reduction in total energy prediction error by a factor of 1347 and an SCF calculation speed-up by a factor of 18%. These substantial improvements set new benchmarks for achieving accurate and applicable predictions in larger molecular systems.

Problem

• Improves Kohn-Sham Hamiltonian scalability for larger molecules
• Introduces Wavefunction Alignment Loss for accurate predictions
• Accelerates self-consistent-field DFT calculations significantly

Innovation

• Deep-learning models predict Hamiltonians
• Wavefunction Alignment Loss improves accuracy
• Larger training set enhances scalability