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

arXiv: https://arxiv.org/abs/2502.19227 · HTML: https://mlsanity.com/paper/2502.19227/

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 96 references that are in the map.
- First posted on arXiv 26 Feb 2025.

## Builds on (most cited first, 30 of 44)

- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Mixtral of Experts](https://mlsanity.com/paper/2401.04088/) (2024)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks](https://mlsanity.com/paper/2006.10503/) (2020)
- [Directional Message Passing for Molecular Graphs](https://mlsanity.com/paper/2003.03123/) (2020)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [ST-MoE: Designing Stable and Transferable Sparse Expert Models](https://mlsanity.com/paper/2202.08906/) (2022)
- [GemNet: Universal Directional Graph Neural Networks for Molecules](https://mlsanity.com/paper/2106.08903/) (2021)
- [Learning from Protein Structure with Geometric Vector Perceptrons](https://mlsanity.com/paper/2009.01411/) (2020)
- [On the Generalization of Equivariance and Convolution in Neural Networks to the Action of Compact Groups](https://mlsanity.com/paper/1802.03690/) (2018)
- [Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs](https://mlsanity.com/paper/2206.11990/) (2022)
- [Cormorant: Covariant Molecular Neural Networks](https://mlsanity.com/paper/1906.04015/) (2019)
- [Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics](https://mlsanity.com/paper/2204.05249/) (2022)
- [Unified Scaling Laws for Routed Language Models](https://mlsanity.com/paper/2202.01169/) (2022)
- [EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations](https://mlsanity.com/paper/2306.12059/) (2023)
- [Spherical Message Passing for 3D Graph Networks](https://mlsanity.com/paper/2102.05013/) (2021)
- [Deep neural network solution of the electronic Schr"odinger equation](https://mlsanity.com/paper/1909.08423/) (2019)
- [Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs](https://mlsanity.com/paper/2302.03655/) (2023)
- [SE(3) Equivariant Graph Neural Networks with Complete Local Frames](https://mlsanity.com/paper/2110.14811/) (2021)
- [Unifying machine learning and quantum chemistry -- a deep neural network for molecular wavefunctions](https://mlsanity.com/paper/1906.10033/) (2019)
- [ComENet: Towards Complete and Efficient Message Passing for 3D Molecular Graphs](https://mlsanity.com/paper/2206.08515/) (2022)
- [A Hitchhiker's Guide to Geometric GNNs for 3D Atomic Systems](https://mlsanity.com/paper/2312.07511/) (2023)
- [SE(3)-equivariant prediction of molecular wavefunctions and electronic densities](https://mlsanity.com/paper/2106.02347/) (2021)
- [TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials](https://mlsanity.com/paper/2306.06482/) (2023)
- [Eigenvalue and Generalized Eigenvalue Problems: Tutorial](https://mlsanity.com/paper/1903.11240/) (2019)

## Built on by (most cited first, 14 of 14)

- [HIP: Hessian Interatomic Potentials without derivatives](https://mlsanity.com/paper/2509.21624/) (2025)
- [High-order Equivariant Flow Matching for Density Functional Theory Hamiltonian Prediction](https://mlsanity.com/paper/2505.18817/) (2025)
- [Learning from the electronic structure of molecules across the periodic table](https://mlsanity.com/paper/2510.00224/) (2025)
- [Efficient Prediction of SO(3)-Equivariant Hamiltonian Matrices via SO(2) Local Frames](https://mlsanity.com/paper/2506.09398/) (2025)
- [Towards A Transferable Acceleration Method for Density Functional Theory](https://mlsanity.com/paper/2509.25724/) (2025)
- [Elign: Equivariant Diffusion Model Alignment from Foundational Machine Learning Force Fields](https://mlsanity.com/paper/2601.21985/) (2026)
- [Transferable SCF-Acceleration through Solver-Aligned Initialization Learning](https://mlsanity.com/paper/2604.21657/) (2026)
- [Advancing Universal Deep Learning for Electronic-Structure Hamiltonian Prediction of Materials](https://mlsanity.com/paper/2509.19877/) (2025)
- [M=oLe-Lambda: Learning the Coupled-Cluster Response State for Energies, Gradients, and Properties](https://mlsanity.com/paper/2605.29622/) (2026)
- [Coupled Cluster con M=oLe: Molecular Orbital Learning for Neural Wavefunctions](https://mlsanity.com/paper/2602.20232/) (2026)
- [A Fixed-Point Neural Operator for Size- and Functional-Transferable Hamiltonian Prediction](https://mlsanity.com/paper/2606.14498/) (2026)
- [MALOQ: Massively Accelerated Learning of Operators for Quantum Transport](https://mlsanity.com/paper/2606.28911/) (2026)
- [Equivariant Flow Matching for Electron Density Prediction](https://mlsanity.com/paper/2610.02651/) (2026)
- [Physics-Aligned Electronic Ground-State Learning Improves Generalization](https://mlsanity.com/paper/2610.10298/) (2026)

## Abstract (arXiv)

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.