Mixture-of-Experts Transformers for Faithfully Deorbitalized Meta-GGA Density Functionals

Published in ChemRxiv, 2025

Graphical Abstract

Abstract

Meta-GGA density functional theory (DFT) is an important method in ab initio materials modelling; however, its computational cost limits applicability for generating large datasets or simulating extended length and time scales, as necessary for modern materials discovery. Deorbitalization is a promising strategy to accelerate meta-GGA DFT by removing the explicit orbital dependence of the exchange-correlation energy and potential using an approximation for the non-local kinetic energy density. Currently, achieving both accuracy and stability in practical deorbitalized simulations remains challenging, as many popular exchange-correlation density functionals are rigidly structured and inherently resistant to deorbitalization. These challenges motivate the need to understand the adaptability of traditional density functionals and to develop more advanced strategies for their design. Guided by this understanding, we adopt a physics-informed deep learning approach to reparameterize popular semi-local meta-GGA density functionals into non-local deorbitalized surrogates that more faithfully mimic their orbital-dependent parent functionals. Our approach is investigated for the deorbitalization of the MS2 exchange density functional of Sun et al. (DOI: 10.1063/1.4789414) using a transformer encoder that leverages multi-head self-attention and automatic differentiation to accurately predict exchange energy densities and partial derivatives across real-space integration grids. The outcome is a deep learning-based procedure for deorbitalizing meta-GGA density functionals that generalizes across diverse chemical bonding environments in both molecules and materials using a sparse mixture-of-experts model. The work shows the potential of advanced deep learning architectures towards constructing accurate surrogate models of expensive electronic structure methods, whilst outlining a path towards robust deorbitalized simulations of molecules and solids.

Citation: A. Chaudhari and A. J. Logsdail, Mixture-of-Experts Transformers for Faithfully Deorbitalized Meta-GGA Density Functionals, ChemRxiv, 2025, DOI: 10.26434/chemrxiv-2025-mrgzj-v2