Accelerated Neural Riemannian Flows for Kohn-Sham Density Functional Theory
Abstract
Density functional theory (DFT) underlies many electronic structure calculations. Typically, a key step in this framework is the iterative self-consistent field (SCF) procedure, which remains computationally expensive. Recent learning-based approaches accelerate parts of this pipeline by improving Hamiltonian prediction or initialization, but still retain the bottleneck at its root. We propose a gauge-symmetric neural Riemannian flows framework for SCF-free electronic structure computation. Our approach directly propagates electronic states on the Grassmann manifold via gauge-equivariant dynamics, rather than relying on a conventional SCF solver. This construction removes the structural bottleneck imposed by the SCF solver and makes the entire propagation trajectory differentiable, suitable for end-to-end learning of the electronic solver. Furthermore, we recast the propagation process at the operator level into a prefix-scan compatible form for efficient parallel execution and acceleration on modern GPUs. Empirically, our method reaches the electronic ground state without SCF, consistently improves energy and orbital related metrics, and achieves speedups of up to 30 times over prior approaches on MD17 molecules. These results point to a new SCF-free direction for electronic structure computation and highlight the potential of learned electronic solvers that are both accelerator-friendly and aligned with end-to-end differentiable learning.
est. 32% chance this paper gets accepted at ICLR 2027.
What do you think this paper will get?
All positions stay anonymous.